From 46b8419bb53c3db5250195fe3b4364c01eb166dd Mon Sep 17 00:00:00 2001 From: Paul Hernandez <60959+phernandez@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:55:33 -0500 Subject: [PATCH 1/8] refactor(core): compile search filters over an explicit ProjectScope (#1562) Signed-off-by: phernandez Co-authored-by: Claude Fable 5.1 --- CHANGELOG.md | 12 + .../repository/note_type_filters.py | 13 +- .../repository/postgres_search_query.py | 667 ++++++++++++++ .../repository/postgres_search_repository.py | 822 +----------------- src/basic_memory/repository/search_filters.py | 147 ++++ .../repository/search_repository_base.py | 20 +- src/basic_memory/repository/search_scope.py | 61 ++ .../repository/sqlite_search_query.py | 423 +++++++++ .../repository/sqlite_search_repository.py | 661 +------------- .../repository/temporal_filters.py | 19 +- tests/repository/test_hybrid_fusion.py | 4 - .../test_postgres_search_quoted_queries.py | 8 +- .../test_postgres_search_repository.py | 85 +- .../test_postgres_search_repository_unit.py | 5 +- .../test_search_file_path_prefix.py | 2 +- .../test_search_relaxed_rendering.py | 16 +- tests/repository/test_search_repository.py | 166 ++-- tests/repository/test_search_scope.py | 47 + tests/repository/test_semantic_search_base.py | 4 - tests/repository/test_semantic_vector_sync.py | 4 - tests/repository/test_vector_pagination.py | 4 - tests/repository/test_vector_threshold.py | 4 - 22 files changed, 1593 insertions(+), 1601 deletions(-) create mode 100644 src/basic_memory/repository/postgres_search_query.py create mode 100644 src/basic_memory/repository/search_filters.py create mode 100644 src/basic_memory/repository/search_scope.py create mode 100644 src/basic_memory/repository/sqlite_search_query.py create mode 100644 tests/repository/test_search_scope.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 239f084dc..4262bb491 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -109,6 +109,18 @@ Frontmatter is now classified once, by the parser, as present, absent, or malformed, and only the first two are ever written to. +### Internal + +- **#1558**: Search filter compilation now runs over an explicit `ProjectScope` instead of + a repository-bound `project_id`. FTS term preparation and filter compilation moved out + of the SQLite and Postgres repositories into `sqlite_search_query` and + `postgres_search_query` as pure functions returning a `CompiledFilter`, and the filters + both backends share (scope, permalink, directory, item type, category, note type, + `after_date`, valid time, candidate keys) are compiled once in `search_filters`. The + note-type and valid-time predicates match search rows on their full + `(project_id, ...)` identity. Project repositories call the compilers with a scope of + one; no query behavior changes. First step of the shared single/multi-project reader. + ## v0.23.2 (2026-08-25) diff --git a/src/basic_memory/repository/note_type_filters.py b/src/basic_memory/repository/note_type_filters.py index 9da86df09..9c8767845 100644 --- a/src/basic_memory/repository/note_type_filters.py +++ b/src/basic_memory/repository/note_type_filters.py @@ -33,6 +33,7 @@ from typing import Any, Sequence +from basic_memory.repository.search_scope import ProjectScope from basic_memory.schemas.search import SearchItemType SEARCH_TABLE = "search_index" @@ -49,6 +50,7 @@ def build_note_type_predicate( note_types: Sequence[str], params: dict[str, Any], *, + scope: ProjectScope, note_type_value: str, ) -> str: """Build the WHERE-clause fragment restricting rows to notes of the given types. @@ -57,19 +59,22 @@ def build_note_type_predicate( documented case-insensitive, so both sides are folded to lowercase. Binds are added to `params` in place, following the convention the surrounding FTS - query builders already use. `project_id` is bound by the caller for the whole query. + query builders already use. The owning note is matched on `(project_id, id)`: the + search row's identity is composite, and the subquery is restricted to `scope` so an + owner outside the caller's projects can never admit a row. """ placeholders = [] for index, note_type in enumerate(note_types): name = f"note_type_{index}" params[name] = note_type.lower() placeholders.append(f":{name}") + owner_scope = scope.predicate(f"{_OWNER}.project_id", params) return ( - f"{SEARCH_TABLE}.entity_id IN (\n" - f" SELECT {_OWNER}.id\n" + f"({SEARCH_TABLE}.project_id, {SEARCH_TABLE}.entity_id) IN (\n" + f" SELECT {_OWNER}.project_id, {_OWNER}.id\n" f" FROM {SEARCH_TABLE} AS {_OWNER}\n" f" WHERE {_OWNER}.type = '{SearchItemType.ENTITY.value}'\n" - f" AND {_OWNER}.project_id = :project_id\n" + f" AND {owner_scope}\n" f" AND LOWER({note_type_value}) IN ({', '.join(placeholders)}))" ) diff --git a/src/basic_memory/repository/postgres_search_query.py b/src/basic_memory/repository/postgres_search_query.py new file mode 100644 index 000000000..ea30070b7 --- /dev/null +++ b/src/basic_memory/repository/postgres_search_query.py @@ -0,0 +1,667 @@ +"""PostgreSQL tsquery preparation: term syntax and filter compilation. + +Pure functions over a ``ProjectScope`` and the caller's filters. Nothing here opens +a session or owns an index. +""" + +import json +import re +from collections.abc import Sequence +from datetime import datetime +from typing import Any + +from basic_memory.repository.metadata_filters import parse_metadata_filters +from basic_memory.repository.script_ngrams import analyze_script_query +from basic_memory.repository.search_filters import ( + AFTER_DATE_ORDER_BY, + POSTGRES_FILTER_DIALECT, + CompiledFilter, + shared_filter_conditions, +) +from basic_memory.repository.search_query import relaxation_word_tokens, relaxed_query_words +from basic_memory.repository.search_repository_base import ( + SearchIndexKey, + metadata_contains_like_condition, + metadata_filter_content_type_condition, +) +from basic_memory.repository.search_scope import ProjectScope +from basic_memory.schemas.search import SearchItemType +from basic_memory.temporal import TemporalFilter + +_TSQUERY_OPERAND_PATTERN = re.compile(r"'(?:''|[^'])*'(?::\*)?|[^\s&|!()]+") +_TSQUERY_WORD_PATTERN = re.compile(r"[^\W_]+(?:'[^\W_]+)?", re.UNICODE) +_QUOTED_QUERY_PATTERN = re.compile(r'"([^"]*)"') +_BOOLEAN_WORDS = frozenset({"AND", "OR", "NOT"}) +_TSQUERY_METACHARACTERS = frozenset("&|!:<>") +# tsquery special characters that must not reach the parser as text. +_TSQUERY_SPECIAL_CHARS = ("&", "|", "!", "(", ")", ":") + + +# --- Term preparation --- + + +def _tsquery_operands(processed_text: str) -> list[tuple[str, str]]: + """Return unique (query operand, representative text) pairs in source order.""" + operands: dict[str, str] = {} + for operand in _TSQUERY_OPERAND_PATTERN.findall(processed_text): + representative = operand.removesuffix(":*") + if representative.startswith("'") and representative.endswith("'"): + representative = representative[1:-1].replace("''", "'") + operands.setdefault(operand, representative) + continue + + # An unquoted apostrophe is invalid tsquery syntax. Keep the literal word for + # the synthetic document, but quote and escape its probe so a strict syntax + # failure can proceed to the relaxed retry. + if "'" in representative: + escaped = "'{}'".format(representative.replace("'", "''")) + safe_operand = f"{escaped}:*" if operand.endswith(":*") else escaped + operands.setdefault(safe_operand, representative) + continue + + # PostgreSQL legitimately parses punctuation inside operands such as + # ``v0.13.0b2:*`` and ``auth-service:*``. Preserve those bytes so the + # synthetic document is tokenized the same way as the original note. + if "<" not in representative and ">" not in representative: + operands.setdefault(operand, representative) + continue + + # A malformed strict operand (for example ``foo str: + """Render user text as complete, individually escaped tsquery operands.""" + words = relaxation_word_tokens(text_value) + if drop_boolean_words: + words = [word for word in words if word.upper() not in _BOOLEAN_WORDS] + if not words: + return "NOSPECIALCHARS:*" + + operands: list[str] = [] + for word in words: + escaped_word = "'{}'".format(word.replace("'", "''")) if "'" in word else word + operands.append(f"{escaped_word}:*" if is_prefix else escaped_word) + return operator.join(operands) + + +def _render_boolean_operand(operand: str) -> str: + """Preserve safe structured text while escaping tsquery syntax bytes.""" + if "'" not in operand and not any(char in _TSQUERY_METACHARACTERS for char in operand): + return operand + return _render_tsquery_words(operand, operator=" & ", is_prefix=False) + + +def _has_valid_boolean_shape(expression: str) -> bool: + """Reject incomplete operator structure before it reaches strict ``to_tsquery``.""" + depth = 0 + for char in expression: + if char == "(": + depth += 1 + elif char == ")": + depth -= 1 + if depth < 0: + return False + if depth: + return False + + stripped = expression.strip() + if not stripped or stripped[0] in "&|" or stripped[-1] in "&|!": + return False + return not any( + re.search(pattern, stripped) + for pattern in ( + r"[&|]\s*[&|]", + r"!\s*[&|)]", + r"\(\s*[&|)]", + r"[&|!(]\s*\)", + ) + ) + + +def prepare_single_term(term: str, is_prefix: bool = True) -> str: + """Prepare one search term with no Boolean operators. + + Multi-word queries become ``word1 & word2``; ``is_prefix`` adds the ``:*`` + suffix; tsquery special characters are removed. + """ + if not term or not term.strip(): + return term + + term = term.strip() + + # An existing wildcard pattern converts to the tsquery prefix operator. + if "*" in term: + return term.replace("*", ":*") + + cleaned_term = term + for char in _TSQUERY_SPECIAL_CHARS: + cleaned_term = cleaned_term.replace(char, " ") + + if " " in cleaned_term: + # Strip sentence punctuation from word edges so question-form queries + # produce clean lexemes (parity with SQLite FTS5 prep). The tsquery tokenizer + # ignores this punctuation anyway; leaving it only risks syntax errors. + words = [w.strip("?!.,;") for w in cleaned_term.split()] + words = [w for w in words if w] + if not words: + # Only special characters remained; emit a term that cannot error. + return "NOSPECIALCHARS:*" + prepared_words = [f"{word}:*" for word in words] if is_prefix else words + return " & ".join(prepared_words) + + # Single word: strip edge punctuation and guard the now-empty case so a bare + # ":*" never reaches tsquery. + cleaned_term = cleaned_term.strip().strip("?!.,;") + if not cleaned_term: + return "NOSPECIALCHARS:*" + if is_prefix: + return f"{cleaned_term}:*" + return cleaned_term + + +def prepare_boolean_query(query: str) -> str: + """Convert a Boolean query to tsquery operators (``&``, ``|``, ``!``). + + Examples: ``coffee AND brewing`` -> ``coffee & brewing``; + ``(pour OR french) AND press`` -> ``(pour | french) & press``; + ``coffee NOT decaf`` -> ``coffee & !decaf``. + """ + # PostgreSQL's strict to_tsquery grammar does not accept web-style double quotes. + # Convert complete quoted groups first so operator-looking words inside them + # remain text and every word becomes a complete operand. + quoted_phrases: dict[str, str] = {} + + def replace_quoted_phrase(match: re.Match[str]) -> str: + phrase = _render_tsquery_words(match.group(1), operator=" & ", is_prefix=False) + placeholder = f"BMQUOTEDPHRASE{len(quoted_phrases)}" + while placeholder in query: + placeholder += "X" + quoted_phrases[placeholder] = f"({phrase})" + # Surround the placeholder so quotes adjacent to plain text become explicit + # operands instead of restoring into ``word(group)``. + return f" {placeholder} " + + result = _QUOTED_QUERY_PATTERN.sub(replace_quoted_phrase, query) + if '"' in result: + # An unmatched quote is user text, not a reason to abort the database + # transaction. Boolean-looking words lose their operator role here. + return _render_tsquery_words(query, operator=" & ", is_prefix=True, drop_boolean_words=True) + + # Boolean syntax is the only structure retained from user input. A + # whitespace-delimited operand still needs explicit conjunctions, but PostgreSQL + # must tokenize structured single operands such as ``auth-service`` and + # ``config.json`` exactly as it did before quoted query normalization. + normalized_parts: list[str] = [] + operator_pattern = r"((? str: + """Prepare user text as a tsquery. + + Boolean operators convert to tsquery form, prefix matching uses ``:*``, and + terms are sanitized so they cannot raise tsquery syntax errors. + """ + if '"' in term or any(op in f" {term} " for op in (" AND ", " OR ", " NOT ")): + return prepare_boolean_query(term) + return prepare_single_term(term, is_prefix) + + +def _relaxed_tsquery_term(word: str) -> str: + """Render one relaxed word as a tsquery-safe prefix expression. + + Mirrors the SQLite renderer: a word token can contain an apostrophe, and tsquery + reads that as lexeme-quoting syntax rather than text. Quoting the lexeme and + doubling any interior quote keeps it literal. + """ + if "'" in word: + return "'{}':*".format(word.replace("'", "''")) + return f"{word}:*" + + +def relaxed_tsquery_text(search_text: str | None) -> str | None: + """OR-relaxed tsquery expression for a failed strict query, or None.""" + words = relaxed_query_words(search_text) + if not words: + return None + return " | ".join(_relaxed_tsquery_term(word) for word in words) + + +def is_tsquery_syntax_error(exc: Exception) -> bool: + msg = str(exc).lower() + return ( + "syntax error in tsquery" in msg + or "invalid input syntax for type tsquery" in msg + or "no operand in tsquery" in msg + or "no operator in tsquery" in msg + ) + + +# --- Filter compilation --- + + +def document_fts_vector_sql(processed_texts: Sequence[str], params: dict[str, Any]) -> str: + """Build a query-sized vector representing lexemes found anywhere in one item.""" + operands: dict[str, str] = {} + for processed_text in processed_texts: + for operand, representative in _tsquery_operands(processed_text): + operands.setdefault(operand, representative) + + present_lexemes: list[str] = [] + for index, (operand, representative) in enumerate(operands.items()): + operand_param = f"text_operand_{index}" + representative_param = f"text_representative_{index}" + params[operand_param] = operand + params[representative_param] = representative + present_lexemes.append( + "CASE WHEN (search_index.textsearchable_index_col " + f"@@ to_tsquery('english', :{operand_param}) OR EXISTS (" + "SELECT 1 FROM search_index_fts_chunks AS operand_chunk " + "WHERE operand_chunk.project_id = search_index.project_id " + "AND operand_chunk.search_index_id = search_index.id " + "AND operand_chunk.search_index_type = search_index.type " + "AND operand_chunk.textsearchable_index_col " + f"@@ to_tsquery('english', :{operand_param}))) " + f"THEN :{representative_param} ELSE '' END" + ) + + if not present_lexemes: + return "search_index.textsearchable_index_col" + + # The synthesized text contains at most the query operands, never the note body. + # This preserves document-wide Boolean semantics without recreating an + # unbounded vector. + lexeme_array = f"ARRAY[{', '.join(present_lexemes)}]" + return f"to_tsvector('english', array_to_string({lexeme_array}, ' '))" + + +def _word_candidate_from_clause(scope: ProjectScope, params: dict[str, Any]) -> str: + """Join search rows to the GIN-indexed candidates for the word channel. + + Trigger: PostgreSQL can extract a required-positive query tree. + Why: OR-ing its operands is a safe indexed superset even when terms live in + different chunks. Pure or optional negation returns ``T`` and must retain all + scoped rows for correct semantics. + Outcome: ordinary and required-positive NOT queries use both GIN indexes; only + genuinely unindexable negation scans the scope. + """ + parent_scope = scope.predicate("candidate_parent.project_id", params) + chunk_scope = scope.predicate("candidate_chunk.project_id", params) + all_scope = scope.predicate("candidate_all.project_id", params) + return f""" + search_index JOIN ( + SELECT + candidate_parent.project_id, + candidate_parent.id, + candidate_parent.type + FROM search_index AS candidate_parent + WHERE {parent_scope} + AND querytree(to_tsquery('english', :text)) <> 'T' + AND candidate_parent.textsearchable_index_col + @@ to_tsquery('english', :text_candidate) + UNION + SELECT + candidate_chunk.project_id, + candidate_chunk.search_index_id AS id, + candidate_chunk.search_index_type AS type + FROM search_index_fts_chunks AS candidate_chunk + WHERE {chunk_scope} + AND querytree(to_tsquery('english', :text)) <> 'T' + AND candidate_chunk.textsearchable_index_col + @@ to_tsquery('english', :text_candidate) + UNION + SELECT + candidate_all.project_id, + candidate_all.id, + candidate_all.type + FROM search_index AS candidate_all + WHERE {all_scope} + AND querytree(to_tsquery('english', :text)) = 'T' + ) AS fts_candidate + ON fts_candidate.project_id = search_index.project_id + AND fts_candidate.id = search_index.id + AND fts_candidate.type = search_index.type + """ + + +def _script_candidate_from_clause(scope: ProjectScope, params: dict[str, Any]) -> str: + """Join search rows to the GIN-indexed candidates for the script channel. + + Trigger: a query contains script grams, with or without word terms. + Why: every script phrase is required, while an English word clause can reduce + to an empty tsquery after dictionary processing. + Outcome: start from the parent and child script GIN indexes, then apply every + word and script predicate. + """ + parent_scope = scope.predicate("script_parent.project_id", params) + chunk_scope = scope.predicate("script_candidate.project_id", params) + return f""" + search_index JOIN ( + SELECT + script_parent.project_id, + script_parent.id, + script_parent.type + FROM search_index AS script_parent + WHERE {parent_scope} + AND script_parent.script_ngrams_index_col + @@ to_tsquery('simple', :script_candidate_text) + UNION + SELECT + script_candidate.project_id, + script_candidate.search_index_id AS id, + script_candidate.search_index_type AS type + FROM search_index_fts_chunks AS script_candidate + WHERE {chunk_scope} + AND script_candidate.script_ngrams_index_col + @@ to_tsquery('simple', :script_candidate_text) + ) AS fts_candidate + ON fts_candidate.project_id = search_index.project_id + AND fts_candidate.id = search_index.id + AND fts_candidate.type = search_index.type + """ + + +def compile_fts_filter( + scope: ProjectScope, + *, + search_text: str | None = None, + permalink: str | None = None, + permalink_match: str | None = None, + title: str | None = None, + note_types: Sequence[str] | None = None, + after_date: datetime | None = None, + search_item_types: Sequence[SearchItemType] | None = None, + categories: Sequence[str] | None = None, + metadata_filters: dict[str, Any] | None = None, + file_path_prefix: str | None = None, + temporal: TemporalFilter | None = None, + allow_relaxed: bool = False, + candidate_keys: Sequence[SearchIndexKey] | None = None, +) -> CompiledFilter: + """Compile Postgres FTS FROM/WHERE/score shared by search and count. + + ``allow_relaxed`` widens the indexed candidate set to the relaxed query's + operands too, so the strict statement and its relaxed retry read the same rows. + """ + params: dict[str, Any] = {} + conditions = shared_filter_conditions( + scope, + params, + dialect=POSTGRES_FILTER_DIALECT, + permalink=permalink, + file_path_prefix=file_path_prefix, + candidate_keys=candidate_keys, + search_item_types=search_item_types, + categories=categories, + note_types=note_types, + after_date=after_date, + temporal=temporal, + ) + from_clause = "search_index" + document_vector: str | None = None + script_tsqueries: list[str] = [] + + # Wildcard-only and blank text add no text condition: every row matches. + if search_text and search_text.strip() not in ("", "*"): + script_query = analyze_script_query(search_text.strip()) + if script_query.word_text: + processed_text = prepare_search_term(script_query.word_text) + params["text"] = processed_text + probe_texts = [processed_text] + if allow_relaxed: + relaxed_text = relaxed_tsquery_text(script_query.word_text) + if relaxed_text: + probe_texts.append(relaxed_text) + + candidate_operands: dict[str, None] = {} + for probe_text in probe_texts: + for operand, _representative in _tsquery_operands(probe_text): + candidate_operands.setdefault(operand, None) + if candidate_operands: + params["text_candidate"] = " | ".join(candidate_operands) + from_clause = _word_candidate_from_clause(scope, params) + document_vector = document_fts_vector_sql(probe_texts, params) + word_condition = f"{document_vector} @@ to_tsquery('english', :text)" + if script_query.gram_phrases: + # Trigger: PostgreSQL's English dictionary removes every word term. + # Why: an empty word query must not suppress a required script match. + # Outcome: only mixed queries treat the empty word channel as neutral; + # word-only stopword queries retain their established empty result. + word_condition = f"(numnode(to_tsquery('english', :text)) = 0 OR {word_condition})" + conditions.append(word_condition) + + if script_query.gram_phrases: + script_tsqueries = [ + " <-> ".join(f"'{gram}'" for gram in phrase) for phrase in script_query.gram_phrases + ] + for index, script_tsquery in enumerate(script_tsqueries): + params[f"script_text_{index}"] = script_tsquery + params["script_candidate_text"] = " | ".join( + f"({script_tsquery})" for script_tsquery in script_tsqueries + ) + from_clause = _script_candidate_from_clause(scope, params) + conditions.extend( + "(search_index.script_ngrams_index_col " + f"@@ to_tsquery('simple', :script_text_{index}) OR EXISTS (" + "SELECT 1 FROM search_index_fts_chunks AS script_chunk " + "WHERE script_chunk.project_id = search_index.project_id " + "AND script_chunk.search_index_id = search_index.id " + "AND script_chunk.search_index_type = search_index.type " + "AND script_chunk.script_ngrams_index_col " + f"@@ to_tsquery('simple', :script_text_{index})))" + for index in range(len(script_tsqueries)) + ) + + if title: + params["title_text"] = prepare_search_term(title.strip(), is_prefix=False) + conditions.append( + "to_tsvector('english', search_index.title) @@ to_tsquery('english', :title_text)" + ) + + if permalink_match: + permalink_text = permalink_match.lower().strip() + if "*" in permalink_match: + # ``*`` becomes the LIKE wildcard. + params["permalink"] = permalink_text.replace("*", "%") + conditions.append("search_index.permalink LIKE :permalink") + else: + params["permalink"] = permalink_text + conditions.append("search_index.permalink = :permalink") + + # Structured metadata filters use jsonb_extract_path_text() / jsonb_extract_path() + # with parameterized path parts instead of #>> / #> with interpolated paths. + if metadata_filters: + parsed_filters = parse_metadata_filters(metadata_filters) + from_clause = f"{from_clause} JOIN entity ON search_index.entity_id = entity.id" + # Frontmatter filters answer for notes only; see + # metadata_filter_content_type_condition for why every regular file would + # otherwise satisfy a null predicate. + conditions.append(metadata_filter_content_type_condition(params)) + metadata_expr = "entity.entity_metadata::jsonb" + + for idx, filt in enumerate(parsed_filters): + path_param_names: list[str] = [] + for j, part in enumerate(filt.path_parts): + path_param = f"meta_path_{idx}_{j}" + params[path_param] = part + path_param_names.append(f":{path_param}") + path_args = ", ".join(path_param_names) + text_expr = f"jsonb_extract_path_text({metadata_expr}, {path_args})" + json_expr = f"jsonb_extract_path({metadata_expr}, {path_args})" + + # jsonb_extract_path_text returns SQL NULL both for a missing key and for + # an explicit JSON null, the same two cases SQLite's json_extract + # collapses, so the dialects answer ``{"owner": None}`` row for row. + # ``= NULL`` is never true. + if filt.op == "is_null": + conditions.append(f"{text_expr} IS NULL") + continue + + if filt.op == "eq": + value_param = f"meta_val_{idx}" + params[value_param] = filt.value + conditions.append(f"{text_expr} = :{value_param}") + continue + + if filt.op == "in": + placeholders: list[str] = [] + for j, val in enumerate(filt.value): + value_param = f"meta_val_{idx}_{j}" + params[value_param] = val + placeholders.append(f":{value_param}") + conditions.append(f"{text_expr} IN ({', '.join(placeholders)})") + continue + + if filt.op == "contains": + base_param = f"meta_val_{idx}" + tag_conditions: list[str] = [] + # Every requested value must be present. + for j, val in enumerate(filt.value): + tag_param = f"{base_param}_{j}" + params[tag_param] = json.dumps([val]) + # The exact JSONB containment test is the primary path; the + # substring patterns only reach values stored as array text. + like_condition = metadata_contains_like_condition( + text_expr, + val, + param_prefix=tag_param, + params=params, + ) + tag_conditions.append( + f"({json_expr} @> CAST(:{tag_param} AS jsonb) OR {like_condition})" + ) + conditions.append(" AND ".join(tag_conditions)) + continue + + if filt.op in {"gt", "gte", "lt", "lte", "between"}: + compare_expr = ( + f"{text_expr}::double precision" if filt.comparison == "numeric" else text_expr + ) + if filt.op == "between": + min_param = f"meta_val_{idx}_min" + max_param = f"meta_val_{idx}_max" + params[min_param] = filt.value[0] + params[max_param] = filt.value[1] + conditions.append(f"{compare_expr} BETWEEN :{min_param} AND :{max_param}") + else: + value_param = f"meta_val_{idx}" + params[value_param] = filt.value + operator = {"gt": ">", "gte": ">=", "lt": "<", "lte": "<="}[filt.op] + conditions.append(f"{compare_expr} {operator} :{value_param}") + continue + + # ts_rank per channel. With no text search there is no rank, so the score is 0. + score_parts: list[str] = [] + if document_vector is not None: + score_parts.append( + "GREATEST(" + f"ts_rank({document_vector}, to_tsquery('english', :text)), " + "ts_rank(search_index.textsearchable_index_col, to_tsquery('english', :text)), " + "COALESCE((SELECT MAX(ts_rank(" + "fts_chunk.textsearchable_index_col, to_tsquery('english', :text))) " + "FROM search_index_fts_chunks AS fts_chunk " + "WHERE fts_chunk.project_id = search_index.project_id " + "AND fts_chunk.search_index_id = search_index.id " + "AND fts_chunk.search_index_type = search_index.type " + "AND fts_chunk.textsearchable_index_col " + "@@ to_tsquery('english', :text)), 0))" + ) + score_parts.extend( + "GREATEST(" + "ts_rank(search_index.script_ngrams_index_col, " + f"to_tsquery('simple', :script_text_{index})), " + "COALESCE((SELECT MAX(ts_rank(script_rank.script_ngrams_index_col, " + f"to_tsquery('simple', :script_text_{index}))) " + "FROM search_index_fts_chunks AS script_rank " + "WHERE script_rank.project_id = search_index.project_id " + "AND script_rank.search_index_id = search_index.id " + "AND script_rank.search_index_type = search_index.type " + "AND script_rank.script_ngrams_index_col " + f"@@ to_tsquery('simple', :script_text_{index})), 0))" + for index in range(len(script_tsqueries)) + ) + # Each condition above is required, so every query component contributes to + # relevance. Taking only the strongest rank would make additional script runs + # invisible. + score_expression = " + ".join(score_parts) if score_parts else "0" + + return CompiledFilter( + from_clause=from_clause, + where_clause=" AND ".join(conditions), + params=params, + order_by_clause=AFTER_DATE_ORDER_BY if after_date else "", + score_expression=score_expression, + ) diff --git a/src/basic_memory/repository/postgres_search_repository.py b/src/basic_memory/repository/postgres_search_repository.py index a6d7e3088..fc89e18da 100644 --- a/src/basic_memory/repository/postgres_search_repository.py +++ b/src/basic_memory/repository/postgres_search_repository.py @@ -2,7 +2,6 @@ import asyncio import json -import re import time from collections.abc import Sequence from datetime import datetime @@ -21,28 +20,23 @@ from basic_memory.repository.rerank_provider import RerankProvider from basic_memory.repository.rerank_provider_factory import create_rerank_provider from basic_memory.repository.search_index_row import SearchIndexRow -from basic_memory.repository.search_query import relaxed_query_words, relaxation_word_tokens -from basic_memory.repository.script_ngrams import analyze_script_query, build_script_ngrams +from basic_memory.repository.search_query import relaxed_query_words +from basic_memory.repository.script_ngrams import build_script_ngrams from basic_memory.repository.semantic_chunking import VectorChunkRecord from basic_memory.repository.search_repository_base import ( SearchIndexKey, SearchRepositoryBase, VectorChunkState, - candidate_key_restriction_condition, - file_path_prefix_condition, - metadata_contains_like_condition, - metadata_filter_content_type_condition, ) from basic_memory.repository.search_trace import ( SearchTraceCollector, build_fts_page_stage, ) -from basic_memory.repository.metadata_filters import parse_metadata_filters -from basic_memory.repository.note_type_filters import ( - POSTGRES_NOTE_TYPE_VALUE, - build_note_type_predicate, +from basic_memory.repository.postgres_search_query import ( + compile_fts_filter, + is_tsquery_syntax_error, + relaxed_tsquery_text, ) -from basic_memory.repository.temporal_filters import build_temporal_predicate from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError from basic_memory.repository.semantic_vector_index import SemanticVectorIndex from basic_memory.repository.semantic_vector_sync import ( @@ -59,111 +53,6 @@ from basic_memory.temporal import TemporalFilter -_TSQUERY_OPERAND_PATTERN = re.compile(r"'(?:''|[^'])*'(?::\*)?|[^\s&|!()]+") -_TSQUERY_WORD_PATTERN = re.compile(r"[^\W_]+(?:'[^\W_]+)?", re.UNICODE) -_QUOTED_QUERY_PATTERN = re.compile(r'"([^"]*)"') -_BOOLEAN_WORDS = frozenset({"AND", "OR", "NOT"}) -_TSQUERY_METACHARACTERS = frozenset("&|!:<>") - - -def _tsquery_operands(processed_text: str) -> list[tuple[str, str]]: - """Return unique (query operand, representative text) pairs in source order.""" - operands: dict[str, str] = {} - for operand in _TSQUERY_OPERAND_PATTERN.findall(processed_text): - representative = operand.removesuffix(":*") - if representative.startswith("'") and representative.endswith("'"): - representative = representative[1:-1].replace("''", "'") - operands.setdefault(operand, representative) - continue - - # An unquoted apostrophe is invalid tsquery syntax. Keep the literal - # word for the synthetic document, but quote and escape its probe so a - # strict syntax failure can proceed to the relaxed retry. - if "'" in representative: - escaped = "'{}'".format(representative.replace("'", "''")) - safe_operand = f"{escaped}:*" if operand.endswith(":*") else escaped - operands.setdefault(safe_operand, representative) - continue - - # PostgreSQL legitimately parses punctuation inside operands such as - # ``v0.13.0b2:*`` and ``auth-service:*``. Preserve those bytes so the - # synthetic document is tokenized the same way as the original note. - if "<" not in representative and ">" not in representative: - operands.setdefault(operand, representative) - continue - - # A malformed strict operand (for example ``foo str: - """Render user text as complete, individually escaped tsquery operands.""" - words = relaxation_word_tokens(text_value) - if drop_boolean_words: - words = [word for word in words if word.upper() not in _BOOLEAN_WORDS] - if not words: - return "NOSPECIALCHARS:*" - - operands = [] - for word in words: - escaped_word = "'{}'".format(word.replace("'", "''")) if "'" in word else word - operands.append(f"{escaped_word}:*" if is_prefix else escaped_word) - return operator.join(operands) - - -def _render_boolean_operand(operand: str) -> str: - """Preserve safe structured text while escaping tsquery syntax bytes.""" - if "'" not in operand and not any(char in _TSQUERY_METACHARACTERS for char in operand): - return operand - return _render_tsquery_words( - operand, - operator=" & ", - is_prefix=False, - ) - - -def _has_valid_boolean_shape(expression: str) -> bool: - """Reject incomplete operator structure before it reaches strict ``to_tsquery``.""" - depth = 0 - for char in expression: - if char == "(": - depth += 1 - elif char == ")": - depth -= 1 - if depth < 0: - return False - if depth: - return False - - stripped = expression.strip() - if not stripped or stripped[0] in "&|" or stripped[-1] in "&|!": - return False - return not any( - re.search(pattern, stripped) - for pattern in ( - r"[&|]\s*[&|]", - r"!\s*[&|)]", - r"\(\s*[&|)]", - r"[&|!(]\s*\)", - ) - ) - - def _strip_nul_from_row(row_data: dict[str, Any]) -> dict[str, Any]: """Strip NUL bytes from all string values in a row dict. @@ -396,236 +285,6 @@ async def _replace_fts_chunks( {"project_id": self.project_id, "chunks": json.dumps(chunks)}, ) - # ------------------------------------------------------------------ - # tsquery preparation (backend-specific) - # ------------------------------------------------------------------ - - @override - def _prepare_search_term(self, term: str, is_prefix: bool = True) -> str: - """Prepare a search term for tsquery format. - - Args: - term: The search term to prepare - is_prefix: Whether to add prefix search capability (:* operator) - - Returns: - Formatted search term for tsquery - - For Postgres: - - Boolean operators are converted to tsquery format (&, |, !) - - Prefix matching uses the :* operator - - Terms are sanitized to prevent tsquery syntax errors - """ - # Check for explicit boolean operators - boolean_operators = [" AND ", " OR ", " NOT "] - if '"' in term or any(op in f" {term} " for op in boolean_operators): - return self._prepare_boolean_query(term) - - # For non-Boolean queries, prepare single term - return self._prepare_single_term(term, is_prefix) - - @staticmethod - def _relaxed_tsquery_term(word: str) -> str: - """Render one relaxed word as a tsquery-safe prefix expression. - - Mirrors the SQLite renderer: a word token can contain an apostrophe, and - tsquery reads that as lexeme-quoting syntax rather than text. Quoting the - lexeme and doubling any interior quote keeps it literal. - """ - if "'" in word: - return "'{}':*".format(word.replace("'", "''")) - return f"{word}:*" - - @staticmethod - def _relaxed_tsquery_text(search_text: Optional[str]) -> Optional[str]: - """OR-relaxed tsquery expression for a failed strict query, or None.""" - words = relaxed_query_words(search_text) - if not words: - return None - return " | ".join(PostgresSearchRepository._relaxed_tsquery_term(word) for word in words) - - def _prepare_boolean_query(self, query: str) -> str: - """Convert Boolean query to tsquery format. - - Args: - query: A Boolean query like "coffee AND brewing" or "(pour OR french) AND press" - - Returns: - tsquery-formatted string with & (AND), | (OR), ! (NOT) operators - - Examples: - "coffee AND brewing" -> "coffee & brewing" - "(pour OR french) AND press" -> "(pour | french) & press" - "coffee NOT decaf" -> "coffee & !decaf" - """ - # PostgreSQL's strict to_tsquery grammar does not accept web-style double - # quotes. Convert complete quoted groups first so operator-looking words - # inside them remain text and every word becomes a complete operand. - quoted_phrases: dict[str, str] = {} - - def replace_quoted_phrase(match: re.Match[str]) -> str: - phrase = _render_tsquery_words( - match.group(1), - operator=" & ", - is_prefix=False, - ) - placeholder = f"BMQUOTEDPHRASE{len(quoted_phrases)}" - while placeholder in query: - placeholder += "X" - quoted_phrases[placeholder] = f"({phrase})" - # Surround the placeholder so quotes adjacent to plain text become - # explicit operands instead of restoring into ``word(group)``. - return f" {placeholder} " - - result = _QUOTED_QUERY_PATTERN.sub(replace_quoted_phrase, query) - if '"' in result: - # An unmatched quote is user text, not a reason to abort the database - # transaction. Boolean-looking words lose their operator role here. - return _render_tsquery_words( - query, - operator=" & ", - is_prefix=True, - drop_boolean_words=True, - ) - - # Boolean syntax is the only structure retained from user input. A - # whitespace-delimited operand still needs explicit conjunctions, but - # PostgreSQL must tokenize structured single operands such as - # ``auth-service`` and ``config.json`` exactly as it did before quoted - # query normalization. - normalized_parts: list[str] = [] - operator_pattern = r"((? str: - """Prepare a single search term for tsquery. - - Args: - term: A single search term - is_prefix: Whether to add prefix search capability (:* suffix) - - Returns: - A properly formatted single term for tsquery - - For Postgres tsquery: - - Multi-word queries become "word1 & word2" - - Prefix matching uses ":*" suffix (e.g., "coff:*") - - Special characters that need escaping: & | ! ( ) : - """ - if not term or not term.strip(): - return term - - term = term.strip() - - # Check if term is already a wildcard pattern - if "*" in term: - # Replace * with :* for Postgres prefix matching - return term.replace("*", ":*") - - # Remove tsquery special characters from the search term - # These characters have special meaning in tsquery and cause syntax errors - # if not used as operators - special_chars = ["&", "|", "!", "(", ")", ":"] - cleaned_term = term - for char in special_chars: - cleaned_term = cleaned_term.replace(char, " ") - - # Handle multi-word queries - if " " in cleaned_term: - # Strip sentence punctuation from word edges so question-form - # queries produce clean lexemes (parity with SQLite FTS5 prep). - # The tsquery tokenizer ignores this punctuation anyway; leaving it - # in only risks tsquery syntax errors. Interior characters are kept. - words = [w.strip("?!.,;") for w in cleaned_term.split()] - words = [w for w in words if w] - if not words: - # All characters were special chars, search won't match anything - # Return a safe search term that won't cause syntax errors - return "NOSPECIALCHARS:*" - if is_prefix: - # Add prefix matching to each word - prepared_words = [f"{word}:*" for word in words] - else: - prepared_words = words - # Join with AND operator - return " & ".join(prepared_words) - - # Single word: strip edge punctuation; guard the now-empty case so a - # bare ":*"/"" never reaches tsquery. - cleaned_term = cleaned_term.strip().strip("?!.,;") - if not cleaned_term: - return "NOSPECIALCHARS:*" - if is_prefix: - return f"{cleaned_term}:*" - else: - return cleaned_term - # ------------------------------------------------------------------ # Abstract hook implementations (vector/semantic, Postgres-specific) # ------------------------------------------------------------------ @@ -957,431 +616,6 @@ async def bulk_index_items(self, search_index_rows: List[SearchIndexRow]) -> Non # FTS search (Postgres-specific) # ------------------------------------------------------------------ - @staticmethod - def _is_tsquery_syntax_error(exc: Exception) -> bool: - msg = str(exc).lower() - return ( - "syntax error in tsquery" in msg - or "invalid input syntax for type tsquery" in msg - or "no operand in tsquery" in msg - or "no operator in tsquery" in msg - ) - - async def _build_fts_query_parts( - self, - search_text: Optional[str] = None, - permalink: Optional[str] = None, - permalink_match: Optional[str] = None, - title: Optional[str] = None, - note_types: Optional[List[str]] = None, - after_date: Optional[datetime] = None, - search_item_types: Optional[List[SearchItemType]] = None, - categories: Optional[List[str]] = None, - metadata_filters: Optional[dict[str, Any]] = None, - file_path_prefix: Optional[str] = None, - temporal: Optional[TemporalFilter] = None, - allow_relaxed: bool = False, - candidate_keys: Sequence[SearchIndexKey] | None = None, - ) -> tuple[str, str, dict[str, Any], str, str]: - """Build Postgres FTS FROM/WHERE params shared by search and count.""" - conditions = [] - params = {} - order_by_clause = "" - from_clause = "search_index" - document_vector_sql: str | None = None - script_tsqueries: list[str] = [] - - # Handle text search for title and content using tsvector - if search_text: - if search_text.strip() == "*" or search_text.strip() == "": - # For wildcard searches, don't add any text conditions - pass - else: - script_query = analyze_script_query(search_text.strip()) - if script_query.word_text: - processed_text = self._prepare_search_term(script_query.word_text) - params["text"] = processed_text - probe_texts = [processed_text] - if allow_relaxed: - relaxed_text = self._relaxed_tsquery_text(script_query.word_text) - if relaxed_text: - probe_texts.append(relaxed_text) - - candidate_operands: dict[str, None] = {} - for probe_text in probe_texts: - for operand, _representative in _tsquery_operands(probe_text): - candidate_operands.setdefault(operand, None) - if candidate_operands: - params["text_candidate"] = " | ".join(candidate_operands) - - # Trigger: PostgreSQL can extract a required-positive query tree. - # Why: OR-ing its operands is a safe indexed superset even when - # terms live in different chunks. Pure/optional negation returns - # ``T`` and must retain all project rows for correct semantics. - # Outcome: ordinary and required-positive NOT queries use both - # GIN indexes; only genuinely unindexable negation scans the project. - from_clause = """ - search_index JOIN ( - SELECT - candidate_parent.project_id, - candidate_parent.id, - candidate_parent.type - FROM search_index AS candidate_parent - WHERE candidate_parent.project_id = :project_id - AND querytree(to_tsquery('english', :text)) <> 'T' - AND candidate_parent.textsearchable_index_col - @@ to_tsquery('english', :text_candidate) - UNION - SELECT - candidate_chunk.project_id, - candidate_chunk.search_index_id AS id, - candidate_chunk.search_index_type AS type - FROM search_index_fts_chunks AS candidate_chunk - WHERE candidate_chunk.project_id = :project_id - AND querytree(to_tsquery('english', :text)) <> 'T' - AND candidate_chunk.textsearchable_index_col - @@ to_tsquery('english', :text_candidate) - UNION - SELECT - candidate_all.project_id, - candidate_all.id, - candidate_all.type - FROM search_index AS candidate_all - WHERE candidate_all.project_id = :project_id - AND querytree(to_tsquery('english', :text)) = 'T' - ) AS fts_candidate - ON fts_candidate.project_id = search_index.project_id - AND fts_candidate.id = search_index.id - AND fts_candidate.type = search_index.type - """ - document_vector_sql = self._document_fts_vector_sql(probe_texts, params) - word_condition = f"{document_vector_sql} @@ to_tsquery('english', :text)" - if script_query.gram_phrases: - # Trigger: PostgreSQL's English dictionary removes every word term. - # Why: an empty word query must not suppress a required script match. - # Outcome: only mixed queries treat the empty word channel as neutral; - # word-only stopword queries retain their established empty result. - word_condition = ( - f"(numnode(to_tsquery('english', :text)) = 0 OR {word_condition})" - ) - conditions.append(word_condition) - - if script_query.gram_phrases: - script_tsqueries = [ - " <-> ".join(f"'{gram}'" for gram in phrase) - for phrase in script_query.gram_phrases - ] - for index, script_tsquery in enumerate(script_tsqueries): - params[f"script_text_{index}"] = script_tsquery - # Trigger: a query contains script grams, with or without word terms. - # Why: every script phrase is required, while an English word clause can - # reduce to an empty tsquery after dictionary processing. - # Outcome: start from the parent and child script GIN indexes, then apply - # every word and script predicate below. - params["script_candidate_text"] = " | ".join( - f"({script_tsquery})" for script_tsquery in script_tsqueries - ) - from_clause = """ - search_index JOIN ( - SELECT - script_parent.project_id, - script_parent.id, - script_parent.type - FROM search_index AS script_parent - WHERE script_parent.project_id = :project_id - AND script_parent.script_ngrams_index_col - @@ to_tsquery('simple', :script_candidate_text) - UNION - SELECT - script_candidate.project_id, - script_candidate.search_index_id AS id, - script_candidate.search_index_type AS type - FROM search_index_fts_chunks AS script_candidate - WHERE script_candidate.project_id = :project_id - AND script_candidate.script_ngrams_index_col - @@ to_tsquery('simple', :script_candidate_text) - ) AS fts_candidate - ON fts_candidate.project_id = search_index.project_id - AND fts_candidate.id = search_index.id - AND fts_candidate.type = search_index.type - """ - conditions.extend( - "(search_index.script_ngrams_index_col " - f"@@ to_tsquery('simple', :script_text_{index}) OR EXISTS (" - "SELECT 1 FROM search_index_fts_chunks AS script_chunk " - "WHERE script_chunk.project_id = search_index.project_id " - "AND script_chunk.search_index_id = search_index.id " - "AND script_chunk.search_index_type = search_index.type " - "AND script_chunk.script_ngrams_index_col " - f"@@ to_tsquery('simple', :script_text_{index})))" - for index in range(len(script_tsqueries)) - ) - - # Handle title search - if title: - title_text = self._prepare_search_term(title.strip(), is_prefix=False) - params["title_text"] = title_text - conditions.append( - "to_tsvector('english', search_index.title) @@ to_tsquery('english', :title_text)" - ) - - # Handle permalink exact search - if permalink: - params["permalink"] = permalink - conditions.append("search_index.permalink = :permalink") - - # Handle permalink pattern match - if permalink_match: - permalink_text = permalink_match.lower().strip() - params["permalink"] = permalink_text - if "*" in permalink_match: - # Use LIKE for pattern matching in Postgres - # Convert * to % for SQL LIKE - permalink_pattern = permalink_text.replace("*", "%") - params["permalink"] = permalink_pattern - conditions.append("search_index.permalink LIKE :permalink") - else: - conditions.append("search_index.permalink = :permalink") - - # Handle directory subtree scope. The predicate is built by the shared - # helper so Postgres and SQLite scope by the identical rule; see - # file_path_prefix_condition for the boundary and escaping reasoning. - subtree_condition = file_path_prefix_condition(file_path_prefix, params) - if subtree_condition is not None: - conditions.append(subtree_condition) - - # Handle an explicit candidate-row restriction. Built by the shared helper so - # both backends restrict by the identical rule; see - # candidate_key_restriction_condition for why the vector filter pass asks about - # its candidates rather than paging the filter's whole match set (#1431). - if candidate_keys is not None: - conditions.append(candidate_key_restriction_condition(candidate_keys, params)) - - # Handle search item type filter (parameterized for defense-in-depth) - if search_item_types: - type_placeholders = [] - for idx, t in enumerate(search_item_types): - param_name = f"search_type_{idx}" - params[param_name] = t.value - type_placeholders.append(f":{param_name}") - conditions.append(f"search_index.type IN ({', '.join(type_placeholders)})") - - # Handle observation category filter (parameterized for defense-in-depth). - # Trigger: caller passed `categories` to scope observation results. - # Why: `entity_types=["observation"]` only narrows to the observation row type; - # callers expect exact-category matching, not incidental text matches. - # Outcome: only rows whose indexed category exactly equals a requested value - # survive (entities/relations have NULL category and are excluded). - if categories: - category_placeholders = [] - for idx, category in enumerate(categories): - param_name = f"category_{idx}" - params[param_name] = category - category_placeholders.append(f":{param_name}") - conditions.append(f"search_index.category IN ({', '.join(category_placeholders)})") - - # Handle note type filter (frontmatter type field, parameterized). - # Trigger: caller passed `note_types` to scope by the frontmatter `type` field. - # Why: the type belongs to the note, but only its entity row carries the - # frontmatter; observation and relation rows do not. Reading it off each row - # silently excluded every non-entity row, which made `note_types` combined - # with a valid-time filter unsatisfiable. - # Outcome: resolved through the owning note in one shared builder, so both - # backends ask the same question and observation rows of a matching note - # are admitted. - if note_types: - conditions.append( - build_note_type_predicate( - note_types, params, note_type_value=POSTGRES_NOTE_TYPE_VALUE - ) - ) - - # Handle date filter - if after_date: - params["after_date"] = after_date - # Filter on updated_at so recently-edited notes are included even when created_at is old - conditions.append("search_index.updated_at > :after_date") - # order by most recent first - order_by_clause = ", search_index.updated_at DESC" - - # Handle authored valid time (SPEC-82). - # Trigger: caller asked when a statement was true of the world. - # Why: `after_date` above filters `updated_at`, which records when the note was - # last edited. That is bookkeeping, never a semantic claim; a decision - # effective through July says nothing about when its file was touched. - # Outcome: an independent predicate over the temporal projection, textually - # identical to the SQLite one because canonical bounds compare - # lexicographically on both backends. Undated sources carry no row and - # are therefore excluded whenever a valid-time filter is present. - if temporal is not None: - conditions.append(build_temporal_predicate(temporal, params)) - - # Handle structured metadata filters (frontmatter) - # Uses jsonb_extract_path_text() / jsonb_extract_path() with parameterized - # path parts instead of #>> / #> with interpolated paths. - if metadata_filters: - parsed_filters = parse_metadata_filters(metadata_filters) - from_clause = f"{from_clause} JOIN entity ON search_index.entity_id = entity.id" - # Frontmatter filters answer for notes only; see - # metadata_filter_content_type_condition for why every regular file - # would otherwise satisfy a null predicate. - conditions.append(metadata_filter_content_type_condition(params)) - metadata_expr = "entity.entity_metadata::jsonb" - - for idx, filt in enumerate(parsed_filters): - # Parameterize each JSON path part individually - path_param_names = [] - for j, part in enumerate(filt.path_parts): - path_param = f"meta_path_{idx}_{j}" - params[path_param] = part - path_param_names.append(f":{path_param}") - path_args = ", ".join(path_param_names) - text_expr = f"jsonb_extract_path_text({metadata_expr}, {path_args})" - json_expr = f"jsonb_extract_path({metadata_expr}, {path_args})" - - # jsonb_extract_path_text returns SQL NULL both for a missing key - # and for an explicit JSON null — the same two cases SQLite's - # json_extract collapses — so the dialects answer - # `{"owner": None}` row for row. `= NULL` is never true, so - # equality here would report a confident zero. - if filt.op == "is_null": - conditions.append(f"{text_expr} IS NULL") - continue - - if filt.op == "eq": - value_param = f"meta_val_{idx}" - params[value_param] = filt.value - conditions.append(f"{text_expr} = :{value_param}") - continue - - if filt.op == "in": - placeholders = [] - for j, val in enumerate(filt.value): - value_param = f"meta_val_{idx}_{j}" - params[value_param] = val - placeholders.append(f":{value_param}") - conditions.append(f"{text_expr} IN ({', '.join(placeholders)})") - continue - - if filt.op == "contains": - base_param = f"meta_val_{idx}" - tag_conditions = [] - # Require all values to be present - for j, val in enumerate(filt.value): - tag_param = f"{base_param}_{j}" - params[tag_param] = json.dumps([val]) - # The exact JSONB containment test is the primary path; the - # substring patterns only reach values stored as array text. - like_condition = metadata_contains_like_condition( - text_expr, - val, - param_prefix=tag_param, - params=params, - ) - tag_conditions.append( - f"({json_expr} @> CAST(:{tag_param} AS jsonb) OR {like_condition})" - ) - conditions.append(" AND ".join(tag_conditions)) - continue - - if filt.op in {"gt", "gte", "lt", "lte", "between"}: - compare_expr = ( - f"{text_expr}::double precision" - if filt.comparison == "numeric" - else text_expr - ) - - if filt.op == "between": - min_param = f"meta_val_{idx}_min" - max_param = f"meta_val_{idx}_max" - params[min_param] = filt.value[0] - params[max_param] = filt.value[1] - conditions.append(f"{compare_expr} BETWEEN :{min_param} AND :{max_param}") - else: - value_param = f"meta_val_{idx}" - params[value_param] = filt.value - operator = {"gt": ">", "gte": ">=", "lt": "<", "lte": "<="}[filt.op] - conditions.append(f"{compare_expr} {operator} :{value_param}") - continue - - # Always filter by project_id - params["project_id"] = self.project_id - conditions.append("search_index.project_id = :project_id") - - # Build WHERE clause - where_clause = " AND ".join(conditions) if conditions else "1=1" - - # Build SQL with ts_rank() for scoring - # Note: If no text search, score will be NULL, so we use COALESCE to default to 0 - score_parts: list[str] = [] - if document_vector_sql is not None: - score_parts.append( - "GREATEST(" - f"ts_rank({document_vector_sql}, to_tsquery('english', :text)), " - "ts_rank(search_index.textsearchable_index_col, to_tsquery('english', :text)), " - "COALESCE((SELECT MAX(ts_rank(" - "fts_chunk.textsearchable_index_col, to_tsquery('english', :text))) " - "FROM search_index_fts_chunks AS fts_chunk " - "WHERE fts_chunk.project_id = search_index.project_id " - "AND fts_chunk.search_index_id = search_index.id " - "AND fts_chunk.search_index_type = search_index.type " - "AND fts_chunk.textsearchable_index_col " - "@@ to_tsquery('english', :text)), 0))" - ) - score_parts.extend( - "GREATEST(" - "ts_rank(search_index.script_ngrams_index_col, " - f"to_tsquery('simple', :script_text_{index})), " - "COALESCE((SELECT MAX(ts_rank(script_rank.script_ngrams_index_col, " - f"to_tsquery('simple', :script_text_{index}))) " - "FROM search_index_fts_chunks AS script_rank " - "WHERE script_rank.project_id = search_index.project_id " - "AND script_rank.search_index_id = search_index.id " - "AND script_rank.search_index_type = search_index.type " - "AND script_rank.script_ngrams_index_col " - f"@@ to_tsquery('simple', :script_text_{index})), 0))" - for index in range(len(script_tsqueries)) - ) - # Each condition above is required, so every query component should contribute to - # relevance. Taking only the strongest rank makes additional script runs invisible. - score_expr = " + ".join(score_parts) if score_parts else "0" - - return from_clause, where_clause, params, order_by_clause, score_expr - - @staticmethod - def _document_fts_vector_sql(processed_texts: Sequence[str], params: dict[str, Any]) -> str: - """Build a query-sized vector representing lexemes found anywhere in one item.""" - operands: dict[str, str] = {} - for processed_text in processed_texts: - for operand, representative in _tsquery_operands(processed_text): - operands.setdefault(operand, representative) - - present_lexemes: list[str] = [] - for index, (operand, representative) in enumerate(operands.items()): - operand_param = f"text_operand_{index}" - representative_param = f"text_representative_{index}" - params[operand_param] = operand - params[representative_param] = representative - present_lexemes.append( - "CASE WHEN (search_index.textsearchable_index_col " - f"@@ to_tsquery('english', :{operand_param}) OR EXISTS (" - "SELECT 1 FROM search_index_fts_chunks AS operand_chunk " - "WHERE operand_chunk.project_id = search_index.project_id " - "AND operand_chunk.search_index_id = search_index.id " - "AND operand_chunk.search_index_type = search_index.type " - "AND operand_chunk.textsearchable_index_col " - f"@@ to_tsquery('english', :{operand_param}))) " - f"THEN :{representative_param} ELSE '' END" - ) - - if not present_lexemes: - return "search_index.textsearchable_index_col" - - # The synthesized text contains at most the query operands, never the note body. - # This preserves document-wide Boolean semantics without recreating an unbounded vector. - lexeme_array = f"ARRAY[{', '.join(present_lexemes)}]" - return f"to_tsvector('english', array_to_string({lexeme_array}, ' '))" - @override async def search( self, @@ -1430,13 +664,8 @@ async def search( return dispatched # --- FTS mode (Postgres-specific) --- - ( - from_clause, - where_clause, - params, - order_by_clause, - score_expr, - ) = await self._build_fts_query_parts( + compiled = compile_fts_filter( + self.scope, search_text=search_text, permalink=permalink, permalink_match=permalink_match, @@ -1451,8 +680,7 @@ async def search( allow_relaxed=allow_relaxed, candidate_keys=candidate_keys, ) - - # set limit and offset + params = compiled.params params["limit"] = limit params["offset"] = offset @@ -1473,10 +701,10 @@ async def search( search_index.category, search_index.created_at, search_index.updated_at, - {score_expr} as score - FROM {from_clause} - WHERE {where_clause} - ORDER BY score DESC {order_by_clause}, search_index.id ASC + {compiled.score_expression} as score + FROM {compiled.from_clause} + WHERE {compiled.where_clause} + ORDER BY score DESC {compiled.order_by_clause}, search_index.id ASC LIMIT :limit OFFSET :offset """ @@ -1498,13 +726,13 @@ async def execute_rows(active_session: AsyncSession, query_params: dict[str, Any return result.fetchall() async def run_search(active_session: AsyncSession): - relaxed = self._relaxed_tsquery_text(search_text) if allow_relaxed else None + relaxed = relaxed_tsquery_text(search_text) if allow_relaxed else None strict_syntax_error = False relaxed_fallback_used = False try: rows = await execute_rows(active_session, params) except Exception as exc: - if not (self._is_tsquery_syntax_error(exc) and relaxed and params.get("text")): + if not (is_tsquery_syntax_error(exc) and relaxed and params.get("text")): raise strict_syntax_error = True rows = [] @@ -1545,7 +773,7 @@ async def run_search(active_session: AsyncSession): async with db.scoped_session(self.session_maker) as owned_session: rows, relaxed_fallback_used = await run_search(owned_session) except Exception as e: - if self._is_tsquery_syntax_error(e): + if is_tsquery_syntax_error(e): logger.warning(f"tsquery syntax error for search term: {search_text}, error: {e}") if trace is not None: trace.fts = build_fts_page_stage( @@ -1619,13 +847,8 @@ async def count( min_similarity=min_similarity, ) - ( - from_clause, - where_clause, - params, - _order_by_clause, - _score_expr, - ) = await self._build_fts_query_parts( + compiled = compile_fts_filter( + self.scope, search_text=search_text, permalink=permalink, permalink_match=permalink_match, @@ -1639,7 +862,8 @@ async def count( temporal=temporal, allow_relaxed=allow_relaxed, ) - sql = f"SELECT COUNT(*) FROM {from_clause} WHERE {where_clause}" + params = compiled.params + sql = f"SELECT COUNT(*) FROM {compiled.from_clause} WHERE {compiled.where_clause}" logger.trace(f"Count {sql} params: {params}") async def execute_count(active_session: AsyncSession, query_params: dict[str, Any]) -> int: @@ -1652,12 +876,12 @@ async def execute_count(active_session: AsyncSession, query_params: dict[str, An try: async with db.scoped_session(self.session_maker) as session: - relaxed = self._relaxed_tsquery_text(search_text) if allow_relaxed else None + relaxed = relaxed_tsquery_text(search_text) if allow_relaxed else None strict_syntax_error = False try: total = await execute_count(session, params) except Exception as exc: - if not (self._is_tsquery_syntax_error(exc) and relaxed and params.get("text")): + if not (is_tsquery_syntax_error(exc) and relaxed and params.get("text")): raise strict_syntax_error = True total = 0 @@ -1675,7 +899,7 @@ async def execute_count(active_session: AsyncSession, query_params: dict[str, An ) return total except Exception as e: - if self._is_tsquery_syntax_error(e): + if is_tsquery_syntax_error(e): logger.warning(f"tsquery syntax error for search term: {search_text}, error: {e}") return 0 logger.error(f"Database error during search count: {e}") diff --git a/src/basic_memory/repository/search_filters.py b/src/basic_memory/repository/search_filters.py new file mode 100644 index 000000000..0a3233bca --- /dev/null +++ b/src/basic_memory/repository/search_filters.py @@ -0,0 +1,147 @@ +"""WHERE-clause pieces both search backends share. + +The two FTS engines differ in how they match text and read JSON. Every other filter a +search accepts asks the same question of the same columns on both, so it is compiled +once here. A backend supplies the two spellings that differ through ``FilterDialect`` +and appends its own text, title, permalink-pattern, and metadata predicates around the +shared ones. +""" + +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import datetime +from typing import Any + +from basic_memory.repository.note_type_filters import ( + POSTGRES_NOTE_TYPE_VALUE, + SQLITE_NOTE_TYPE_VALUE, + build_note_type_predicate, +) +from basic_memory.repository.search_repository_base import ( + SearchIndexKey, + candidate_key_restriction_condition, + file_path_prefix_condition, +) +from basic_memory.repository.search_scope import ProjectScope +from basic_memory.repository.temporal_filters import build_temporal_predicate +from basic_memory.schemas.search import SearchItemType +from basic_memory.temporal import TemporalFilter + +# Newest edits first whenever the caller filtered on ``after_date``. +AFTER_DATE_ORDER_BY = ", search_index.updated_at DESC" + + +@dataclass(frozen=True, slots=True) +class FilterDialect: + """The two SQL spellings that differ between backends inside the shared filters.""" + + note_type_value: str + after_date_condition: str + + +SQLITE_FILTER_DIALECT = FilterDialect( + note_type_value=SQLITE_NOTE_TYPE_VALUE, + # datetime() normalizes both sides so ISO strings of mixed precision compare as instants. + after_date_condition="datetime(search_index.updated_at) > datetime(:after_date)", +) +POSTGRES_FILTER_DIALECT = FilterDialect( + note_type_value=POSTGRES_NOTE_TYPE_VALUE, + after_date_condition="search_index.updated_at > :after_date", +) + + +@dataclass(frozen=True, slots=True) +class CompiledFilter: + """One backend's FROM, WHERE, and score for a search, ready to place in a statement. + + ``params`` is the bind dictionary the statement runs with. Callers add ``limit`` + and ``offset``, and a relaxed retry replaces ``text``. + """ + + from_clause: str + where_clause: str + params: dict[str, Any] + order_by_clause: str + score_expression: str + + +def shared_filter_conditions( + scope: ProjectScope, + params: dict[str, Any], + *, + dialect: FilterDialect, + permalink: str | None, + file_path_prefix: str | None, + candidate_keys: Sequence[SearchIndexKey] | None, + search_item_types: Sequence[SearchItemType] | None, + categories: Sequence[str] | None, + note_types: Sequence[str] | None, + after_date: datetime | None, + temporal: TemporalFilter | None, +) -> list[str]: + """Compile the filters whose SQL is identical on both backends. + + Binds are added to ``params`` in place. The scope predicate comes first: it is the + one filter every statement carries, and it is what keeps rows outside the caller's + projects out of every candidate window. + """ + conditions = [scope.predicate("search_index.project_id", params)] + + if permalink: + params["permalink"] = permalink + conditions.append("search_index.permalink = :permalink") + + # See file_path_prefix_condition for the subtree boundary and escaping rules. + subtree_condition = file_path_prefix_condition(file_path_prefix, params) + if subtree_condition is not None: + conditions.append(subtree_condition) + + # See candidate_key_restriction_condition for why the vector filter pass asks + # about its candidates rather than paging the filter's whole match set (#1431). + if candidate_keys is not None: + conditions.append(candidate_key_restriction_condition(candidate_keys, params)) + + if search_item_types: + type_placeholders: list[str] = [] + for index, item_type in enumerate(search_item_types): + name = f"search_type_{index}" + params[name] = item_type.value + type_placeholders.append(f":{name}") + conditions.append(f"search_index.type IN ({', '.join(type_placeholders)})") + + # Trigger: caller passed ``categories`` to scope observation results. + # Why: ``entity_types=["observation"]`` only narrows to the observation row type; + # callers expect exact-category matching, not incidental text matches. + # Outcome: only rows whose indexed category exactly equals a requested value + # survive (entities/relations have NULL category and are excluded). + if categories: + category_placeholders: list[str] = [] + for index, category in enumerate(categories): + name = f"category_{index}" + params[name] = category + category_placeholders.append(f":{name}") + conditions.append(f"search_index.category IN ({', '.join(category_placeholders)})") + + # The note type belongs to the note, but only its entity row carries the + # frontmatter, so the predicate resolves through the owning note. See + # note_type_filters for why reading it off each row excluded every non-entity row. + if note_types: + conditions.append( + build_note_type_predicate( + note_types, params, scope=scope, note_type_value=dialect.note_type_value + ) + ) + + # Filter on updated_at so recently edited notes are included even when created_at + # is old. The matching ORDER BY lives in AFTER_DATE_ORDER_BY. + if after_date: + params["after_date"] = after_date + conditions.append(dialect.after_date_condition) + + # Authored valid time (SPEC-82) is independent of ``after_date``: that one is + # bookkeeping about the file, this one is a claim about the world. See + # temporal_filters for the overlap rule and why the subquery is non-correlated. + if temporal is not None: + conditions.append(build_temporal_predicate(temporal, params, scope=scope)) + + return conditions diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 7e824a86e..99f8b23ae 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -33,6 +33,7 @@ validate_rerank_scores, ) from basic_memory.repository.search_index_row import SearchIndexRow +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.script_ngrams import build_script_ngrams from basic_memory.repository.search_trace import ( BelowThreshold, @@ -405,6 +406,8 @@ def __init__(self, session_maker: async_sessionmaker[AsyncSession], project_id: self.session_maker = session_maker self.project_id = project_id + # Every statement this repository compiles reads exactly one project. + self.scope = ProjectScope.single(project_id) async def semantic_effectively_enabled(self) -> bool: """Return whether semantic retrieval can actually run for this repository. @@ -455,23 +458,6 @@ async def init_search_index(self) -> None: """ pass - @abstractmethod - def _prepare_search_term(self, term: str, is_prefix: bool = True) -> str: - """Prepare a search term for backend-specific query syntax. - - Args: - term: The search term to prepare - is_prefix: Whether to add prefix search capability - - Returns: - Formatted search term for the backend - - Backend-specific implementations: - - SQLite: Quotes FTS5 special characters, adds * wildcards - - Postgres: Converts to tsquery syntax with :* prefix operator - """ - pass - @abstractmethod async def search( self, diff --git a/src/basic_memory/repository/search_scope.py b/src/basic_memory/repository/search_scope.py new file mode 100644 index 000000000..c6fa78137 --- /dev/null +++ b/src/basic_memory/repository/search_scope.py @@ -0,0 +1,61 @@ +"""The project set a search statement is allowed to read. + +Every search row, vector chunk, and temporal assertion carries a ``project_id``. A +statement binds that column to an explicit set before any ranking runs, so rows +outside the set never occupy a candidate window. Absence is not a value here: a scope +is always built from concrete IDs, and an empty scope compiles to a predicate that +admits nothing. +""" + +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Any + +# No project can match, and no bind needs to be sent to prove it. +_MATCHES_NOTHING = "1 = 0" + + +@dataclass(frozen=True, slots=True) +class ProjectScope: + """Unique, positive project IDs a statement may read, in ascending order. + + Build one with ``ProjectScope.of`` or ``ProjectScope.single``; both canonicalize + the input so two scopes over the same projects compare equal. + """ + + project_ids: tuple[int, ...] + + def __post_init__(self) -> None: + for project_id in self.project_ids: + # bool is an int subclass; ``True`` would silently read as project 1. + if isinstance(project_id, bool) or not isinstance(project_id, int) or project_id <= 0: + raise ValueError(f"Project IDs must be positive integers, got {project_id!r}") + + @classmethod + def of(cls, project_ids: Iterable[int]) -> "ProjectScope": + """Canonicalize any iterable of project IDs into a scope.""" + return cls(tuple(sorted(set(project_ids)))) + + @classmethod + def single(cls, project_id: int) -> "ProjectScope": + """The scope every project-bound repository runs under.""" + return cls((project_id,)) + + @property + def is_empty(self) -> bool: + return not self.project_ids + + def predicate(self, column: str, params: dict[str, Any]) -> str: + """SQL restricting ``column`` to this scope, adding its binds to ``params``. + + The bind names are a function of the scope alone, so a statement that + references the scope from several subqueries sends each ID once. + """ + if self.is_empty: + return _MATCHES_NOTHING + placeholders: list[str] = [] + for index, project_id in enumerate(self.project_ids): + name = f"scope_{index}" + params[name] = project_id + placeholders.append(f":{name}") + return f"{column} IN ({', '.join(placeholders)})" diff --git a/src/basic_memory/repository/sqlite_search_query.py b/src/basic_memory/repository/sqlite_search_query.py new file mode 100644 index 000000000..9fa161bab --- /dev/null +++ b/src/basic_memory/repository/sqlite_search_query.py @@ -0,0 +1,423 @@ +"""SQLite FTS5 query preparation: term syntax and filter compilation. + +Pure functions over a ``ProjectScope`` and the caller's filters. Nothing here opens +a session or owns an index; the repository resolves the one piece of live schema +state the compiler needs (the entity table's columns) and passes it in. +""" + +import re +from collections.abc import Collection, Sequence +from datetime import datetime +from typing import Any + +from basic_memory.repository.metadata_filters import build_sqlite_json_path, parse_metadata_filters +from basic_memory.repository.script_ngrams import analyze_script_query +from basic_memory.repository.search_filters import ( + AFTER_DATE_ORDER_BY, + SQLITE_FILTER_DIALECT, + CompiledFilter, + shared_filter_conditions, +) +from basic_memory.repository.search_query import relaxed_query_words +from basic_memory.repository.search_repository_base import ( + SearchIndexKey, + metadata_contains_like_condition, + metadata_filter_content_type_condition, +) +from basic_memory.repository.search_scope import ProjectScope +from basic_memory.schemas.search import SearchItemType +from basic_memory.temporal import TemporalFilter + +SQLITE_WORD_COLUMNS = "{title content_stems content_snippet}" + +# Characters that indicate a term should be quoted (parentheses excluded: valid syntax). +_NEEDS_QUOTING_CHARS = frozenset(" .:;,<>?/-'\"[]{}+!@#$%^&=|\\~`") +# Characters that can cause FTS5 syntax errors when read as operators. +_PROBLEMATIC_CHARS = frozenset("\"'()[]{}+!@#$%^&=|\\~`") +# Characters that indicate quoting for spaces, dots, colons, and hyphens followed by +# wildcards, which FTS5 mishandles. +_SPACE_OR_SPECIAL_CHARS = frozenset(" .:;,<>?/-") +_BOOLEAN_OPERATOR_PATTERN = r"(\bAND\b|\bOR\b|\bNOT\b)" + + +# --- Term preparation --- + + +def needs_quoting(term: str) -> bool: + """Whether a term must be quoted for FTS5 safety.""" + if not term or not term.strip(): + return False + return any(c in _NEEDS_QUOTING_CHARS for c in term) + + +def prepare_single_term(term: str, is_prefix: bool = True) -> str: + """Prepare one search term with no Boolean operators. + + ``is_prefix`` adds the ``*`` suffix so simple terms match by prefix. + """ + if not term or not term.strip(): + return term + + term = term.strip() + + # A proper wildcard pattern ("hello*", "test*world") is left alone. + if "*" in term and all(c.isalnum() or c in "*_-" for c in term): + return term + + # Natural-language queries arrive with sentence punctuation that FTS5 treats as + # syntax ("When did Melanie paint a sunrise?"). The tokenizer ignores this + # punctuation in the index, so stripping it from word edges loses nothing, but + # leaving it forces the whole question into an exact-phrase match that returns + # zero rows and silently disables the FTS half of hybrid search. Interior + # characters (hyphens, slashes in permalinks and paths) are untouched. + if " " in term: + words = [word.strip("?!.,;:") for word in term.split()] + term = " ".join(word for word in words if word) + if not term: + return "" + + has_problematic = any(c in _PROBLEMATIC_CHARS for c in term) + has_spaces_or_special = any(c in _SPACE_OR_SPECIAL_CHARS for c in term) + + if has_problematic or has_spaces_or_special: + if " " in term and not has_problematic: + words = term.split() + has_special_in_words = any( + any(c in word for c in _SPACE_OR_SPECIAL_CHARS if c != " ") for word in words + ) + if not has_special_in_words: + # Multi-word queries of simple words ("emoji unicode") use Boolean AND + # so word order does not matter. + prepared_words = [f"{word}*" for word in words] if is_prefix else words + return " AND ".join(prepared_words) + # Any word with special characters quotes the entire phrase. + escaped_term = term.replace('"', '""') + if is_prefix and not ("/" in term and term.endswith(".md")): + return f'"{escaped_term}"*' + return f'"{escaped_term}"' # pragma: no cover + + # Terms with problematic characters or file paths use exact phrase matching. + escaped_term = term.replace('"', '""') + if is_prefix and not ("/" in term and term.endswith(".md")): + return f'"{escaped_term}"*' + return f'"{escaped_term}"' + + if is_prefix: + return f"{term}*" + return term + + +def prepare_parenthetical_term(term: str) -> str: + """Prepare a term containing parentheses, preserving them for grouping.""" + result = "" + index = 0 + while index < len(term): + if term[index] in "()": + result += term[index] + index += 1 + continue + start = index + while index < len(term) and term[index] not in "()": + index += 1 + content = term[start:index].strip() + if content: + # Quote only when the content needs it; simple words stay bare. + if needs_quoting(content): + escaped_content = content.replace('"', '""') + result += f'"{escaped_content}"' + else: + result += content + return result + + +def prepare_boolean_query(query: str) -> str: + """Quote the terms of a Boolean query while preserving its operators and grouping.""" + processed_parts: list[str] = [] + for part in re.split(_BOOLEAN_OPERATOR_PATTERN, query): + part = part.strip() + if not part: + continue + if part in ("AND", "OR", "NOT"): + processed_parts.append(part) + elif "(" in part or ")" in part: + processed_parts.append(prepare_parenthetical_term(part)) + else: + # Boolean queries do not get prefix wildcards. + processed_parts.append(prepare_single_term(part, is_prefix=False)) + return " ".join(processed_parts) + + +def prepare_search_term(term: str, is_prefix: bool = True) -> str: + """Prepare user text as an FTS5 query. + + Boolean operators (AND, OR, NOT) are preserved. Terms with FTS5 special + characters are quoted. Simple terms get prefix wildcards. + """ + if any(op in f" {term} " for op in (" AND ", " OR ", " NOT ")): + return prepare_boolean_query(term) + return prepare_single_term(term, is_prefix) + + +def _relaxed_fts_term(word: str) -> str: + """Render one relaxed word as an FTS5-safe prefix expression. + + A word token can contain an apostrophe ("об'єкт", "don't"). Interpolated bare it + is FTS5 syntax, not text: the whole expression fails to parse, the caller swallows + the syntax error, and the relaxed retry returns nothing, which is the silent-empty + FTS failure this fallback exists to prevent. + """ + if "'" in word or '"' in word: + return '"{}"*'.format(word.replace('"', '""')) + return f"{word}*" + + +def relaxed_fts_text(search_text: str | None) -> str | None: + """OR-relaxed FTS5 expression for a failed strict query, or None.""" + words = relaxed_query_words(search_text) + if not words: + return None + return " OR ".join(_relaxed_fts_term(word) for word in words) + + +def is_fts5_syntax_error(exc: Exception) -> bool: + return "fts5: syntax error" in str(exc).lower() + + +# --- Filter compilation --- + + +def compile_fts_filter( + scope: ProjectScope, + *, + entity_columns: Collection[str], + search_text: str | None = None, + permalink: str | None = None, + permalink_match: str | None = None, + title: str | None = None, + note_types: Sequence[str] | None = None, + after_date: datetime | None = None, + search_item_types: Sequence[SearchItemType] | None = None, + categories: Sequence[str] | None = None, + metadata_filters: dict[str, Any] | None = None, + file_path_prefix: str | None = None, + temporal: TemporalFilter | None = None, + candidate_keys: Sequence[SearchIndexKey] | None = None, +) -> CompiledFilter: + """Compile SQLite FTS FROM/WHERE/score shared by search and count. + + ``entity_columns`` is the live column set of the ``entity`` table. Generated + frontmatter columns are used when present and fall back to ``json_extract``. + """ + params: dict[str, Any] = {} + conditions = shared_filter_conditions( + scope, + params, + dialect=SQLITE_FILTER_DIALECT, + permalink=permalink, + file_path_prefix=file_path_prefix, + candidate_keys=candidate_keys, + search_item_types=search_item_types, + categories=categories, + note_types=note_types, + after_date=after_date, + temporal=temporal, + ) + match_conditions: list[str] = [] + from_clause = "search_index" + score_expression = "bm25(search_index)" + preserve_match_score = False + + # Wildcard-only and blank text add no text condition: every row matches. + if search_text and search_text.strip() not in ("", "*"): + script_query = analyze_script_query(search_text.strip()) + # Trigger: the query contains text from an unsegmented script. + # Why: the script channel needs one table-level MATCH alongside word fields. + # Outcome: mixed queries rank all terms together; word-only queries retain + # their established per-column matching and ranking behavior. + if script_query.gram_phrases: + preserve_match_score = True + params["text"] = "" + params["script_text"] = "" + if script_query.word_text: + prepared_text = prepare_search_term(script_query.word_text) + params["text"] = ( + f"(title: ({prepared_text}) OR " + f"content_stems: ({prepared_text}) OR " + f"content_snippet: ({prepared_text}))" + ) + script_phrases = " AND ".join( + f'"{" ".join(phrase)}"' for phrase in script_query.gram_phrases + ) + script_clause = f"script_ngrams: ({script_phrases})" + params["script_text"] = ( + f" AND ({script_clause})" if script_query.word_text else script_clause + ) + match_conditions.append("search_index MATCH (:text || :script_text)") + else: + word_text = ( + script_query.word_text + if script_query.word_text is not None + else search_text.strip() + ) + params["text"] = prepare_search_term(word_text) + # content_stems is capped for Postgres index-row compatibility, while + # SQLite stores the complete note body in its FTS5 content_snippet column. + match_conditions.append( + "(search_index.title MATCH :text OR " + "search_index.content_stems MATCH :text OR " + "search_index.content_snippet MATCH :text)" + ) + + if title: + params["title_text"] = prepare_search_term(title.strip(), is_prefix=False) + match_conditions.append("search_index.title MATCH :title_text") + + if permalink_match: + # GLOB patterns keep their syntax; prepare_search_term would quote the slashes. + permalink_text = permalink_match.lower().strip() + params["permalink"] = permalink_text + if "*" in permalink_match: + conditions.append("search_index.permalink GLOB :permalink") + elif "/" in permalink_text: + conditions.append("search_index.permalink = :permalink") + else: + # A bare name without a path matches through FTS5. + params["permalink"] = prepare_search_term(permalink_text, is_prefix=False) + match_conditions.append("search_index.permalink MATCH :permalink") + + if metadata_filters: + parsed_filters = parse_metadata_filters(metadata_filters) + from_clause = "search_index JOIN entity ON search_index.entity_id = entity.id" + # Frontmatter filters answer for notes only; see + # metadata_filter_content_type_condition for why every regular file would + # otherwise satisfy a null predicate. + conditions.append(metadata_filter_content_type_condition(params)) + + for idx, filt in enumerate(parsed_filters): + path_param = f"meta_path_{idx}" + extract_expr = None + use_tags_column = False + + if filt.path_parts == ["status"] and "frontmatter_status" in entity_columns: + extract_expr = "entity.frontmatter_status" + elif filt.path_parts == ["type"] and "frontmatter_type" in entity_columns: + extract_expr = "entity.frontmatter_type" + elif filt.path_parts == ["tags"] and "tags_json" in entity_columns: + extract_expr = "entity.tags_json" + use_tags_column = True + + if extract_expr is None: + params[path_param] = build_sqlite_json_path(filt.path_parts) + extract_expr = f"json_extract(entity.entity_metadata, :{path_param})" + + # json_extract returns SQL NULL both for a missing key and for an explicit + # JSON null, and the generated frontmatter_* columns are that same + # json_extract, so IS NULL means "the note carries no value here", the + # question ``{"owner": None}`` asks. ``= NULL`` is never true. + if filt.op == "is_null": + conditions.append(f"{extract_expr} IS NULL") + continue + + if filt.op == "eq": + value_param = f"meta_val_{idx}" + params[value_param] = filt.value + conditions.append(f"{extract_expr} = :{value_param}") + continue + + if filt.op == "in": + placeholders: list[str] = [] + for j, val in enumerate(filt.value): + value_param = f"meta_val_{idx}_{j}" + params[value_param] = val + placeholders.append(f":{value_param}") + conditions.append(f"{extract_expr} IN ({', '.join(placeholders)})") + continue + + if filt.op == "contains": + tag_conditions: list[str] = [] + for j, val in enumerate(filt.value): + value_param = f"meta_val_{idx}_{j}" + params[value_param] = val + # The exact JSON-membership test is the primary path; the + # substring patterns only reach values stored as array text. + like_condition = metadata_contains_like_condition( + extract_expr, + val, + param_prefix=value_param, + params=params, + ) + json_each_expr = ( + "json_each(entity.tags_json)" + if use_tags_column + else f"json_each(entity.entity_metadata, :{path_param})" + ) + tag_conditions.append( + "(" + f"EXISTS (SELECT 1 FROM {json_each_expr} WHERE value = :{value_param}) " + f"OR {like_condition}" + ")" + ) + conditions.append(" AND ".join(tag_conditions)) + continue + + if filt.op in {"gt", "gte", "lt", "lte", "between"}: + compare_expr = ( + f"CAST({extract_expr} AS REAL)" + if filt.comparison == "numeric" + else extract_expr + ) + if filt.op == "between": + min_param = f"meta_val_{idx}_min" + max_param = f"meta_val_{idx}_max" + params[min_param] = filt.value[0] + params[max_param] = filt.value[1] + conditions.append(f"{compare_expr} BETWEEN :{min_param} AND :{max_param}") + else: + value_param = f"meta_val_{idx}" + params[value_param] = filt.value + operator = {"gt": ">", "gte": ">=", "lt": "<", "lte": "<="}[filt.op] + conditions.append(f"{compare_expr} {operator} :{value_param}") + continue + + # Trigger: SQLite rejects some Boolean combinations of MATCH predicates, + # including a word-field OR expression combined with the script channel. + # Why: each MATCH must be evaluated in an FTS-valid query context. + # Outcome: keep one outer MATCH for bm25 ranking and intersect the rest by rowid. + if len(match_conditions) > 1: + ranked_match, *additional_matches = match_conditions + conditions.extend( + f"search_index.rowid IN (SELECT rowid FROM search_index WHERE {match_condition})" + for match_condition in additional_matches + ) + match_conditions = [ranked_match] + + # Trigger: SQLite FTS MATCH predicates combined with JOINs can fail with + # "unable to use function MATCH in the requested context". + # Why: script queries need MATCH and bm25 together for ranking, while legacy + # word-column OR predicates cannot evaluate bm25 in the same derived query. + # Outcome: rank script matches before joining metadata; retain the established + # rowid-filter path for word-only searches. + if metadata_filters and match_conditions: + match_where = " AND ".join(match_conditions) + if preserve_match_score: + from_clause = ( + "(SELECT search_index.rowid AS rowid, search_index.*, " + "bm25(search_index) AS fts_score " + f"FROM search_index WHERE {match_where}) AS search_index " + "JOIN entity ON search_index.entity_id = entity.id" + ) + score_expression = "search_index.fts_score" + else: + conditions.append( + f"search_index.rowid IN (SELECT rowid FROM search_index WHERE {match_where})" + ) + else: + conditions.extend(match_conditions) + + return CompiledFilter( + from_clause=from_clause, + where_clause=" AND ".join(conditions), + params=params, + order_by_clause=AFTER_DATE_ORDER_BY if after_date else "", + score_expression=score_expression, + ) diff --git a/src/basic_memory/repository/sqlite_search_repository.py b/src/basic_memory/repository/sqlite_search_repository.py index 4160d7932..11bb39ff0 100644 --- a/src/basic_memory/repository/sqlite_search_repository.py +++ b/src/basic_memory/repository/sqlite_search_repository.py @@ -1,7 +1,6 @@ """SQLite FTS5-based search repository implementation.""" import asyncio -import re import time from collections.abc import Sequence from contextlib import asynccontextmanager @@ -31,22 +30,18 @@ from basic_memory.repository.search_repository_base import ( SearchIndexKey, SearchRepositoryBase, - candidate_key_restriction_condition, - file_path_prefix_condition, - metadata_contains_like_condition, - metadata_filter_content_type_condition, ) from basic_memory.repository.script_ngrams import analyze_script_query from basic_memory.repository.search_trace import ( SearchTraceCollector, build_fts_page_stage, ) -from basic_memory.repository.metadata_filters import parse_metadata_filters, build_sqlite_json_path -from basic_memory.repository.note_type_filters import ( - SQLITE_NOTE_TYPE_VALUE, - build_note_type_predicate, +from basic_memory.repository.sqlite_search_query import ( + SQLITE_WORD_COLUMNS, + compile_fts_filter, + is_fts5_syntax_error, + relaxed_fts_text, ) -from basic_memory.repository.temporal_filters import build_temporal_predicate from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError from basic_memory.repository.semantic_vector_index import SemanticVectorIndex from basic_memory.repository.semantic_vector_sync import StagedVectorDeletion @@ -56,9 +51,6 @@ from basic_memory.temporal import TemporalFilter -SQLITE_WORD_COLUMNS = "{title content_stems content_snippet}" - - class SQLiteSearchRepository(SearchRepositoryBase): """SQLite FTS5 implementation of search repository. @@ -160,281 +152,6 @@ async def init_search_index(self): ) self._semantic_enabled = False - # ------------------------------------------------------------------ - # FTS5 query preparation (backend-specific) - # ------------------------------------------------------------------ - - def _prepare_boolean_query(self, query: str) -> str: - """Prepare a Boolean query by quoting individual terms while preserving operators. - - Args: - query: A Boolean query like "tier1-test AND unicode" or "(hello OR world) NOT test" - - Returns: - A properly formatted Boolean query with quoted terms that need quoting - """ - # Define Boolean operators and their boundaries - boolean_pattern = r"(\bAND\b|\bOR\b|\bNOT\b)" - - # Split the query by Boolean operators, keeping the operators - parts = re.split(boolean_pattern, query) - - processed_parts = [] - for part in parts: - part = part.strip() - if not part: - continue - - # If it's a Boolean operator, keep it as is - if part in ["AND", "OR", "NOT"]: - processed_parts.append(part) - else: - # Handle parentheses specially - they should be preserved for grouping - if "(" in part or ")" in part: - # Parse parenthetical expressions carefully - processed_part = self._prepare_parenthetical_term(part) - processed_parts.append(processed_part) - else: - # This is a search term - for Boolean queries, don't add prefix wildcards - prepared_term = self._prepare_single_term(part, is_prefix=False) - processed_parts.append(prepared_term) - - return " ".join(processed_parts) - - def _prepare_parenthetical_term(self, term: str) -> str: - """Prepare a term that contains parentheses, preserving the parentheses for grouping. - - Args: - term: A term that may contain parentheses like "(hello" or "world)" or "(hello OR world)" - - Returns: - A properly formatted term with parentheses preserved - """ - # Handle terms that start/end with parentheses but may contain quotable content - result = "" - i = 0 - while i < len(term): - if term[i] in "()": - # Preserve parentheses as-is - result += term[i] - i += 1 - else: - # Find the next parenthesis or end of string - start = i - while i < len(term) and term[i] not in "()": - i += 1 - - # Extract the content between parentheses - content = term[start:i].strip() - if content: - # Only quote if it actually needs quoting (has hyphens, special chars, etc) - # but don't quote if it's just simple words - if self._needs_quoting(content): - escaped_content = content.replace('"', '""') - result += f'"{escaped_content}"' - else: - result += content - - return result - - def _needs_quoting(self, term: str) -> bool: - """Check if a term needs to be quoted for FTS5 safety. - - Args: - term: The term to check - - Returns: - True if the term should be quoted - """ - if not term or not term.strip(): - return False - - # Characters that indicate we should quote (excluding parentheses which are valid syntax) - needs_quoting_chars = [ - " ", - ".", - ":", - ";", - ",", - "<", - ">", - "?", - "/", - "-", - "'", - '"', - "[", - "]", - "{", - "}", - "+", - "!", - "@", - "#", - "$", - "%", - "^", - "&", - "=", - "|", - "\\", - "~", - "`", - ] - - return any(c in term for c in needs_quoting_chars) - - def _prepare_single_term(self, term: str, is_prefix: bool = True) -> str: - """Prepare a single search term (no Boolean operators). - - Args: - term: A single search term - is_prefix: Whether to add prefix search capability (* suffix) - - Returns: - A properly formatted single term - """ - if not term or not term.strip(): - return term - - term = term.strip() - - # Check if term is already a proper wildcard pattern (alphanumeric + *) - # e.g., "hello*", "test*world" - these should be left alone - if "*" in term and all(c.isalnum() or c in "*_-" for c in term): - return term - - # Natural-language queries arrive with sentence punctuation that FTS5 - # treats as syntax ("When did Melanie paint a sunrise?"). The tokenizer - # ignores this punctuation in the INDEX, so stripping it from word - # edges loses nothing — but leaving it forces the whole question into - # an exact-phrase match that returns zero rows, silently disabling the - # FTS half of hybrid search. Interior characters (hyphens, slashes — - # permalinks and paths) are untouched. - if " " in term: - words = [word.strip("?!.,;:") for word in term.split()] - term = " ".join(word for word in words if word) - if not term: - return "" - - # Characters that can cause FTS5 syntax errors when used as operators - # We're more conservative here - only quote when we detect problematic patterns - problematic_chars = [ - '"', - "'", - "(", - ")", - "[", - "]", - "{", - "}", - "+", - "!", - "@", - "#", - "$", - "%", - "^", - "&", - "=", - "|", - "\\", - "~", - "`", - ] - - # Characters that indicate we should quote (spaces, dots, colons, etc.) - # Adding hyphens here because FTS5 can have issues with hyphens followed by wildcards - needs_quoting_chars = [" ", ".", ":", ";", ",", "<", ">", "?", "/", "-"] - - # Check if term needs quoting - has_problematic = any(c in term for c in problematic_chars) - has_spaces_or_special = any(c in term for c in needs_quoting_chars) - - if has_problematic or has_spaces_or_special: - # Handle multi-word queries differently from special character queries - if " " in term and not any(c in term for c in problematic_chars): - # Check if any individual word contains special characters that need quoting - words = term.strip().split() - has_special_in_words = any( - any(c in word for c in needs_quoting_chars if c != " ") for word in words - ) - - if not has_special_in_words: - # For multi-word queries with simple words (like "emoji unicode"), - # use boolean AND to handle word order variations - if is_prefix: - # Add prefix wildcard to each word for better matching - prepared_words = [f"{word}*" for word in words if word] - else: - prepared_words = words - term = " AND ".join(prepared_words) - else: - # If any word has special characters, quote the entire phrase - escaped_term = term.replace('"', '""') - if is_prefix and not ("/" in term and term.endswith(".md")): - term = f'"{escaped_term}"*' - else: - term = f'"{escaped_term}"' # pragma: no cover - else: - # For terms with problematic characters or file paths, use exact phrase matching - # Escape any existing quotes by doubling them - escaped_term = term.replace('"', '""') - # Quote the entire term to handle special characters safely - if is_prefix and not ("/" in term and term.endswith(".md")): - # For search terms (not file paths), add prefix matching - term = f'"{escaped_term}"*' - else: - # For file paths, use exact matching - term = f'"{escaped_term}"' - elif is_prefix: - # Only add wildcard for simple terms without special characters - term = f"{term}*" - - return term - - @override - def _prepare_search_term(self, term: str, is_prefix: bool = True) -> str: - """Prepare a search term for FTS5 query. - - Args: - term: The search term to prepare - is_prefix: Whether to add prefix search capability (* suffix) - - For FTS5: - - Boolean operators (AND, OR, NOT) are preserved for complex queries - - Terms with FTS5 special characters are quoted to prevent syntax errors - - Simple terms get prefix wildcards for better matching - """ - # Check for explicit boolean operators - if present, process as Boolean query - boolean_operators = [" AND ", " OR ", " NOT "] - if any(op in f" {term} " for op in boolean_operators): - return self._prepare_boolean_query(term) - - # For non-Boolean queries, use the single term preparation logic - return self._prepare_single_term(term, is_prefix) - - @staticmethod - def _relaxed_fts_term(word: str) -> str: - """Render one relaxed word as an FTS5-safe prefix expression. - - A word token can contain an apostrophe ("об'єкт", "don't"). Interpolated - bare it is FTS5 syntax, not text: the whole expression fails to parse, the - caller swallows the syntax error, and the relaxed retry returns nothing — - the exact silent-empty-FTS failure this fallback exists to prevent. - """ - if "'" in word or '"' in word: - return '"{}"*'.format(word.replace('"', '""')) - return f"{word}*" - - @staticmethod - def _relaxed_fts_text(search_text: Optional[str]) -> Optional[str]: - """OR-relaxed FTS5 expression for a failed strict query, or None.""" - words = relaxed_query_words(search_text) - if not words: - return None - return " OR ".join(SQLiteSearchRepository._relaxed_fts_term(word) for word in words) - @override async def semantic_effectively_enabled(self) -> bool: """Probe the sqlite-vec runtime instead of trusting still-enabled config. @@ -780,325 +497,6 @@ async def bulk_index_items(self, search_index_rows: List[SearchIndexRow]) -> Non # FTS search (backend-specific) # ------------------------------------------------------------------ - @staticmethod - def _is_fts5_syntax_error(exc: Exception) -> bool: - return "fts5: syntax error" in str(exc).lower() - - async def _build_fts_query_parts( - self, - search_text: Optional[str] = None, - permalink: Optional[str] = None, - permalink_match: Optional[str] = None, - title: Optional[str] = None, - note_types: Optional[List[str]] = None, - after_date: Optional[datetime] = None, - search_item_types: Optional[List[SearchItemType]] = None, - categories: Optional[List[str]] = None, - metadata_filters: Optional[dict[str, Any]] = None, - file_path_prefix: Optional[str] = None, - temporal: Optional[TemporalFilter] = None, - candidate_keys: Sequence[SearchIndexKey] | None = None, - ) -> tuple[str, str, dict[str, Any], str, str]: - """Build SQLite FTS FROM/WHERE params shared by search and count.""" - conditions = [] - match_conditions = [] - params = {} - order_by_clause = "" - from_clause = "search_index" - score_expression = "bm25(search_index)" - preserve_match_score = False - - # Handle text search for title and content - if search_text: - # Skip FTS for wildcard-only queries that would cause "unknown special query" errors - if search_text.strip() == "*" or search_text.strip() == "": - # For wildcard searches, don't add any text conditions - return all results - pass - else: - script_query = analyze_script_query(search_text.strip()) - # Trigger: the query contains text from an unsegmented script. - # Why: the script channel needs one table-level MATCH alongside word fields. - # Outcome: mixed queries rank all terms together; word-only queries retain their - # established per-column matching and ranking behavior. - if script_query.gram_phrases: - preserve_match_score = True - params["text"] = "" - params["script_text"] = "" - if script_query.word_text: - prepared_text = self._prepare_search_term(script_query.word_text) - params["text"] = ( - f"(title: ({prepared_text}) OR " - f"content_stems: ({prepared_text}) OR " - f"content_snippet: ({prepared_text}))" - ) - script_phrases = " AND ".join( - f'"{" ".join(phrase)}"' for phrase in script_query.gram_phrases - ) - script_clause = f"script_ngrams: ({script_phrases})" - params["script_text"] = ( - f" AND ({script_clause})" if script_query.word_text else script_clause - ) - match_conditions.append("search_index MATCH (:text || :script_text)") - else: - word_text = ( - script_query.word_text - if script_query.word_text is not None - else search_text.strip() - ) - processed_text = self._prepare_search_term(word_text) - params["text"] = processed_text - # content_stems is capped for Postgres index-row compatibility, while - # SQLite stores the complete note body in its FTS5 content_snippet column. - match_conditions.append( - "(search_index.title MATCH :text OR " - "search_index.content_stems MATCH :text OR " - "search_index.content_snippet MATCH :text)" - ) - - # Handle title match search - if title: - title_text = self._prepare_search_term(title.strip(), is_prefix=False) - params["title_text"] = title_text - match_conditions.append("search_index.title MATCH :title_text") - - # Handle permalink exact search - if permalink: - params["permalink"] = permalink - conditions.append("search_index.permalink = :permalink") - - # Handle permalink match search, supports * - if permalink_match: - # For GLOB patterns, don't use _prepare_search_term as it will quote slashes - # GLOB patterns need to preserve their syntax - permalink_text = permalink_match.lower().strip() - params["permalink"] = permalink_text - if "*" in permalink_match: - conditions.append("search_index.permalink GLOB :permalink") - else: - # For exact matches without *, we can use FTS5 MATCH - # but only prepare the term if it doesn't look like a path - if "/" in permalink_text: - conditions.append("search_index.permalink = :permalink") - else: - permalink_text = self._prepare_search_term(permalink_text, is_prefix=False) - params["permalink"] = permalink_text - match_conditions.append("search_index.permalink MATCH :permalink") - - # Handle directory subtree scope. The predicate is built by the shared - # helper so SQLite and Postgres scope by the identical rule; see - # file_path_prefix_condition for the boundary and escaping reasoning. - subtree_condition = file_path_prefix_condition(file_path_prefix, params) - if subtree_condition is not None: - conditions.append(subtree_condition) - - # Handle an explicit candidate-row restriction. Built by the shared helper so - # both backends restrict by the identical rule; see - # candidate_key_restriction_condition for why the vector filter pass asks about - # its candidates rather than paging the filter's whole match set (#1431). - if candidate_keys is not None: - conditions.append(candidate_key_restriction_condition(candidate_keys, params)) - - # Handle entity type filter (parameterized for defense-in-depth) - if search_item_types: - type_placeholders = [] - for idx, t in enumerate(search_item_types): - param_name = f"search_type_{idx}" - params[param_name] = t.value - type_placeholders.append(f":{param_name}") - conditions.append(f"search_index.type IN ({', '.join(type_placeholders)})") - - # Handle observation category filter (parameterized for defense-in-depth). - # Trigger: caller passed `categories` to scope observation results. - # Why: `entity_types=["observation"]` only narrows to the observation row type; - # callers expect exact-category matching, not incidental text matches. - # Outcome: only rows whose indexed category exactly equals a requested value - # survive (entities/relations have NULL category and are excluded). - if categories: - category_placeholders = [] - for idx, category in enumerate(categories): - param_name = f"category_{idx}" - params[param_name] = category - category_placeholders.append(f":{param_name}") - conditions.append(f"search_index.category IN ({', '.join(category_placeholders)})") - - # Handle note type filter (frontmatter type field, parameterized). - # Trigger: caller passed `note_types` to scope by the frontmatter `type` field. - # Why: the type belongs to the note, but only its entity row carries the - # frontmatter; observation and relation rows do not. Reading it off each row - # silently excluded every non-entity row, which made `note_types` combined - # with a valid-time filter unsatisfiable. - # Outcome: resolved through the owning note in one shared builder, so both - # backends ask the same question and observation rows of a matching note - # are admitted. - if note_types: - conditions.append( - build_note_type_predicate( - note_types, params, note_type_value=SQLITE_NOTE_TYPE_VALUE - ) - ) - - # Handle date filter using datetime() for proper comparison - if after_date: - params["after_date"] = after_date - # Filter on updated_at so recently-edited notes are included even when created_at is old - conditions.append("datetime(search_index.updated_at) > datetime(:after_date)") - - # order by most recent first - order_by_clause = ", search_index.updated_at DESC" - - # Handle authored valid time (SPEC-82). - # Trigger: caller asked when a statement was true of the world. - # Why: `after_date` above filters `updated_at`, which records when the note was - # last edited. That is bookkeeping, never a semantic claim; a decision - # effective through July says nothing about when its file was touched. - # Outcome: an independent predicate over the temporal projection. It matches - # only sources carrying a structured qualifier, so undated sources are - # excluded whenever a valid-time filter is present, and no ordering - # changes -- relevance still decides the ranking. - if temporal is not None: - conditions.append(build_temporal_predicate(temporal, params)) - - # Handle structured metadata filters (frontmatter) - if metadata_filters: - parsed_filters = parse_metadata_filters(metadata_filters) - from_clause = "search_index JOIN entity ON search_index.entity_id = entity.id" - # Frontmatter filters answer for notes only; see - # metadata_filter_content_type_condition for why every regular file - # would otherwise satisfy a null predicate. - conditions.append(metadata_filter_content_type_condition(params)) - entity_columns = await self._get_entity_columns() - - for idx, filt in enumerate(parsed_filters): - path_param = f"meta_path_{idx}" - extract_expr = None - use_tags_column = False - - if filt.path_parts == ["status"] and "frontmatter_status" in entity_columns: - extract_expr = "entity.frontmatter_status" - elif filt.path_parts == ["type"] and "frontmatter_type" in entity_columns: - extract_expr = "entity.frontmatter_type" - elif filt.path_parts == ["tags"] and "tags_json" in entity_columns: - extract_expr = "entity.tags_json" - use_tags_column = True - - if extract_expr is None: - params[path_param] = build_sqlite_json_path(filt.path_parts) - extract_expr = f"json_extract(entity.entity_metadata, :{path_param})" - - # json_extract returns SQL NULL both for a missing key and for an - # explicit JSON null, and the generated frontmatter_* columns are - # that same json_extract — so IS NULL means "the note carries no - # value here", the question `{"owner": None}` asks. `= NULL` is - # never true, so equality here would report a confident zero. - if filt.op == "is_null": - conditions.append(f"{extract_expr} IS NULL") - continue - - if filt.op == "eq": - value_param = f"meta_val_{idx}" - params[value_param] = filt.value - conditions.append(f"{extract_expr} = :{value_param}") - continue - - if filt.op == "in": - placeholders = [] - for j, val in enumerate(filt.value): - value_param = f"meta_val_{idx}_{j}" - params[value_param] = val - placeholders.append(f":{value_param}") - conditions.append(f"{extract_expr} IN ({', '.join(placeholders)})") - continue - - if filt.op == "contains": - tag_conditions = [] - for j, val in enumerate(filt.value): - value_param = f"meta_val_{idx}_{j}" - params[value_param] = val - # The exact JSON-membership test is the primary path; the - # substring patterns only reach values stored as array text. - like_condition = metadata_contains_like_condition( - extract_expr, - val, - param_prefix=value_param, - params=params, - ) - json_each_expr = ( - "json_each(entity.tags_json)" - if use_tags_column - else f"json_each(entity.entity_metadata, :{path_param})" - ) - tag_conditions.append( - "(" - f"EXISTS (SELECT 1 FROM {json_each_expr} WHERE value = :{value_param}) " - f"OR {like_condition}" - ")" - ) - conditions.append(" AND ".join(tag_conditions)) - continue - - if filt.op in {"gt", "gte", "lt", "lte", "between"}: - compare_expr = ( - f"CAST({extract_expr} AS REAL)" - if filt.comparison == "numeric" - else extract_expr - ) - - if filt.op == "between": - min_param = f"meta_val_{idx}_min" - max_param = f"meta_val_{idx}_max" - params[min_param] = filt.value[0] - params[max_param] = filt.value[1] - conditions.append(f"{compare_expr} BETWEEN :{min_param} AND :{max_param}") - else: - value_param = f"meta_val_{idx}" - params[value_param] = filt.value - operator = {"gt": ">", "gte": ">=", "lt": "<", "lte": "<="}[filt.op] - conditions.append(f"{compare_expr} {operator} :{value_param}") - continue - - # Trigger: SQLite rejects some Boolean combinations of MATCH predicates, - # including a word-field OR expression combined with the script channel. - # Why: each MATCH must be evaluated in an FTS-valid query context. - # Outcome: keep one outer MATCH for bm25 ranking and intersect the rest by rowid. - if len(match_conditions) > 1: - ranked_match, *additional_matches = match_conditions - conditions.extend( - f"search_index.rowid IN (SELECT rowid FROM search_index WHERE {match_condition})" - for match_condition in additional_matches - ) - match_conditions = [ranked_match] - - # Trigger: SQLite FTS MATCH predicates combined with JOINs can fail with - # "unable to use function MATCH in the requested context". - # Why: script queries need MATCH and bm25 together for ranking, while legacy - # word-column OR predicates cannot evaluate bm25 in the same derived query. - # Outcome: rank script matches before joining metadata; retain the established - # rowid-filter path for word-only searches. - if metadata_filters and match_conditions: - match_where = " AND ".join(match_conditions) - if preserve_match_score: - from_clause = ( - "(SELECT search_index.rowid AS rowid, search_index.*, " - "bm25(search_index) AS fts_score " - f"FROM search_index WHERE {match_where}) AS search_index " - "JOIN entity ON search_index.entity_id = entity.id" - ) - score_expression = "search_index.fts_score" - else: - conditions.append( - f"search_index.rowid IN (SELECT rowid FROM search_index WHERE {match_where})" - ) - else: - conditions.extend(match_conditions) - - # Always filter by project_id - params["project_id"] = self.project_id - conditions.append("search_index.project_id = :project_id") - - # Build WHERE clause - where_clause = " AND ".join(conditions) if conditions else "1=1" - return from_clause, where_clause, params, order_by_clause, score_expression - @override async def search( self, @@ -1153,13 +551,11 @@ async def search( return dispatched # --- FTS mode (SQLite-specific) --- - ( - from_clause, - where_clause, - params, - order_by_clause, - score_expression, - ) = await self._build_fts_query_parts( + # Generated frontmatter columns are read only when a metadata filter needs them. + entity_columns = await self._get_entity_columns() if metadata_filters else frozenset() + compiled = compile_fts_filter( + self.scope, + entity_columns=entity_columns, search_text=search_text, permalink=permalink, permalink_match=permalink_match, @@ -1173,8 +569,7 @@ async def search( temporal=temporal, candidate_keys=candidate_keys, ) - - # set limit on search query + params = compiled.params params["limit"] = limit params["offset"] = offset relaxed_search_text = search_text @@ -1198,10 +593,10 @@ async def search( search_index.category, search_index.created_at, search_index.updated_at, - {score_expression} as score - FROM {from_clause} - WHERE {where_clause} - ORDER BY score ASC {order_by_clause} + {compiled.score_expression} as score + FROM {compiled.from_clause} + WHERE {compiled.where_clause} + ORDER BY score ASC {compiled.order_by_clause} LIMIT :limit OFFSET :offset """ @@ -1221,9 +616,7 @@ async def run_search(active_session: AsyncSession): # vector-only. # Outcome: one retry with OR-joined prefix terms; bm25 still # ranks multi-term matches first. - relaxed = ( - self._relaxed_fts_text(relaxed_search_text) if allow_relaxed and not rows else None - ) + relaxed = relaxed_fts_text(relaxed_search_text) if allow_relaxed and not rows else None if relaxed and params.get("text"): relaxed_fallback_used = True params["text"] = ( @@ -1252,7 +645,7 @@ async def run_search(active_session: AsyncSession): rows, relaxed_fallback_used = await run_search(owned_session) except Exception as e: # Handle FTS5 syntax errors and provide user-friendly feedback - if self._is_fts5_syntax_error(e): # pragma: no cover + if is_fts5_syntax_error(e): # pragma: no cover logger.warning(f"FTS5 syntax error for search term: {search_text}, error: {e}") # Return empty results rather than crashing if trace is not None: @@ -1327,13 +720,10 @@ async def count( min_similarity=min_similarity, ) - ( - from_clause, - where_clause, - params, - _order_by_clause, - _score_expression, - ) = await self._build_fts_query_parts( + entity_columns = await self._get_entity_columns() if metadata_filters else frozenset() + compiled = compile_fts_filter( + self.scope, + entity_columns=entity_columns, search_text=search_text, permalink=permalink, permalink_match=permalink_match, @@ -1346,7 +736,8 @@ async def count( file_path_prefix=file_path_prefix, temporal=temporal, ) - sql = f"SELECT COUNT(*) FROM {from_clause} WHERE {where_clause}" + params = compiled.params + sql = f"SELECT COUNT(*) FROM {compiled.from_clause} WHERE {compiled.where_clause}" logger.trace(f"Count {sql} params: {params}") relaxed_search_text = search_text if search_text and "script_text" in params: @@ -1356,9 +747,7 @@ async def count( result = await session.execute(text(sql), params) total = int(result.scalar_one()) relaxed = ( - self._relaxed_fts_text(relaxed_search_text) - if allow_relaxed and total == 0 - else None + relaxed_fts_text(relaxed_search_text) if allow_relaxed and total == 0 else None ) if relaxed and params.get("text"): params["text"] = ( @@ -1375,7 +764,7 @@ async def count( total = int(result.scalar_one()) return total except Exception as e: - if self._is_fts5_syntax_error(e): # pragma: no cover + if is_fts5_syntax_error(e): # pragma: no cover logger.warning(f"FTS5 syntax error for search term: {search_text}, error: {e}") return 0 logger.error(f"Database error during search count: {e}") diff --git a/src/basic_memory/repository/temporal_filters.py b/src/basic_memory/repository/temporal_filters.py index 81c3c666c..d49802122 100644 --- a/src/basic_memory/repository/temporal_filters.py +++ b/src/basic_memory/repository/temporal_filters.py @@ -33,6 +33,7 @@ from typing import Any +from basic_memory.repository.search_scope import ProjectScope from basic_memory.temporal import TemporalFilter, TemporalRange TEMPORAL_INDEX_TABLE = "memory_time_index" @@ -84,7 +85,9 @@ def _not_window_ends_before_source(window: TemporalRange) -> str | None: return f"({' OR '.join(clauses)})" -def build_temporal_predicate(temporal: TemporalFilter, params: dict[str, Any]) -> str: +def build_temporal_predicate( + temporal: TemporalFilter, params: dict[str, Any], *, scope: ProjectScope +) -> str: """Build the WHERE-clause fragment restricting search rows by authored valid time. Two intervals overlap exactly when neither lies entirely before the other, which @@ -97,7 +100,8 @@ def build_temporal_predicate(temporal: TemporalFilter, params: dict[str, Any]) - documented default for a valid-time query. Binds are added to `params` in place, following the convention already used by the - surrounding FTS query builders. + surrounding FTS query builders. Assertions are read from `scope` only, and the + search row is matched on its full `(project_id, type, id)` identity. """ window = temporal.window if window is not None and window.is_empty: @@ -105,7 +109,7 @@ def build_temporal_predicate(temporal: TemporalFilter, params: dict[str, Any]) - # false constant is both correct and cheaper than running the subquery. return _MATCHES_NOTHING - conditions = [f"{TEMPORAL_INDEX_TABLE}.project_id = :project_id"] + conditions = [scope.predicate(f"{TEMPORAL_INDEX_TABLE}.project_id", params)] if temporal.kind is not None: params["tq_kind"] = temporal.kind.value @@ -135,11 +139,12 @@ def build_temporal_predicate(temporal: TemporalFilter, params: dict[str, Any]) - ) where_clause = "\n AND ".join(conditions) - # (type, id) is the search row's own identity and the address this projection - # stores, so the pair joins the two without a correlated reference. + # (project_id, type, id) is the search row's own identity and the address this + # projection stores, so the triple joins the two without a correlated reference. return ( - "(search_index.type, search_index.id) IN (\n" - f" SELECT {TEMPORAL_INDEX_TABLE}.source_type, {TEMPORAL_INDEX_TABLE}.source_id\n" + "(search_index.project_id, search_index.type, search_index.id) IN (\n" + f" SELECT {TEMPORAL_INDEX_TABLE}.project_id, " + f"{TEMPORAL_INDEX_TABLE}.source_type, {TEMPORAL_INDEX_TABLE}.source_id\n" f" FROM {TEMPORAL_INDEX_TABLE}\n" f" WHERE {where_clause})" ) diff --git a/tests/repository/test_hybrid_fusion.py b/tests/repository/test_hybrid_fusion.py index 9e49048f7..4bf7135b1 100644 --- a/tests/repository/test_hybrid_fusion.py +++ b/tests/repository/test_hybrid_fusion.py @@ -71,10 +71,6 @@ async def init_search_index(self): async def get_entity_physical_chunk_keys(self, entity_id: int) -> set[str] | None: return None # physical storage is not inspectable in this double - @override - def _prepare_search_term(self, term, is_prefix=True): - return term # pragma: no cover - @override async def search( self, diff --git a/tests/repository/test_postgres_search_quoted_queries.py b/tests/repository/test_postgres_search_quoted_queries.py index a42f709f5..40a412619 100644 --- a/tests/repository/test_postgres_search_quoted_queries.py +++ b/tests/repository/test_postgres_search_quoted_queries.py @@ -5,6 +5,7 @@ import pytest +import basic_memory.repository.postgres_search_repository as postgres_search_repository_module from basic_memory.repository.postgres_search_repository import PostgresSearchRepository from basic_memory.repository.search_index_row import SearchIndexRow @@ -60,7 +61,7 @@ async def test_quoted_or_phrases_complete_without_tsquery_recovery( ) syntax_errors: list[Exception] = [] - real_is_syntax_error = repository._is_tsquery_syntax_error + real_is_syntax_error = postgres_search_repository_module.is_tsquery_syntax_error def record_syntax_error(exception: Exception) -> bool: is_syntax_error = real_is_syntax_error(exception) @@ -68,7 +69,10 @@ def record_syntax_error(exception: Exception) -> bool: syntax_errors.append(exception) return is_syntax_error - monkeypatch.setattr(repository, "_is_tsquery_syntax_error", record_syntax_error) + # The repository module binds the classifier at import; patch it where it is read. + monkeypatch.setattr( + postgres_search_repository_module, "is_tsquery_syntax_error", record_syntax_error + ) query = '"incident response" OR "database recovery"' async with asyncio.timeout(2): diff --git a/tests/repository/test_postgres_search_repository.py b/tests/repository/test_postgres_search_repository.py index e25b785e8..308509ac7 100644 --- a/tests/repository/test_postgres_search_repository.py +++ b/tests/repository/test_postgres_search_repository.py @@ -13,6 +13,14 @@ from basic_memory.config import BasicMemoryConfig, DatabaseBackend import basic_memory.repository.search_repository_base as search_repository_base_module from basic_memory.repository.litellm_provider import LiteLLMEmbeddingProvider +import basic_memory.repository.postgres_search_query as postgres_search_query_module +import basic_memory.repository.postgres_search_repository as postgres_search_repository_module +from basic_memory.repository.postgres_search_query import ( + compile_fts_filter, + prepare_search_term, + prepare_single_term, + relaxed_tsquery_text, +) from basic_memory.repository.postgres_search_repository import ( PostgresSearchRepository, _strip_nul_from_row, @@ -261,35 +269,29 @@ async def test_postgres_search_repository_bulk_index_items_and_prepare_terms( await repo.bulk_index_items([]) # Exercise term preparation helpers - assert "&" in repo._prepare_search_term("coffee AND brewing") - assert repo._prepare_search_term("coff*") == "coff:*" - assert repo._prepare_search_term("()&!:") == "NOSPECIALCHARS:*" - assert repo._prepare_search_term("coffee brewing") == "coffee:* & brewing:*" - assert repo._prepare_single_term(" ") == " " - assert repo._prepare_single_term("coffee", is_prefix=False) == "coffee" - - indexed_from, _where, indexed_params, _order, _score = await repo._build_fts_query_parts( - search_text="coffee brewing", - allow_relaxed=True, - ) - assert "FROM search_index AS candidate_parent" in indexed_from - assert "FROM search_index_fts_chunks AS candidate_chunk" in indexed_from - assert "querytree(to_tsquery('english', :text))" in indexed_from - assert indexed_params["text_candidate"] == "coffee:* | brewing:*" - - filtered_from, _where, _params, _order, _score = await repo._build_fts_query_parts( - search_text="coffee brewing", - metadata_filters={"status": "active"}, + assert "&" in prepare_search_term("coffee AND brewing") + assert prepare_search_term("coff*") == "coff:*" + assert prepare_search_term("()&!:") == "NOSPECIALCHARS:*" + assert prepare_search_term("coffee brewing") == "coffee:* & brewing:*" + assert prepare_single_term(" ") == " " + assert prepare_single_term("coffee", is_prefix=False) == "coffee" + + indexed = compile_fts_filter(repo.scope, search_text="coffee brewing", allow_relaxed=True) + assert "FROM search_index AS candidate_parent" in indexed.from_clause + assert "FROM search_index_fts_chunks AS candidate_chunk" in indexed.from_clause + assert "querytree(to_tsquery('english', :text))" in indexed.from_clause + assert indexed.params["text_candidate"] == "coffee:* | brewing:*" + + filtered = compile_fts_filter( + repo.scope, search_text="coffee brewing", metadata_filters={"status": "active"} ) - assert "AS fts_candidate" in filtered_from - assert "JOIN entity ON search_index.entity_id = entity.id" in filtered_from + assert "AS fts_candidate" in filtered.from_clause + assert "JOIN entity ON search_index.entity_id = entity.id" in filtered.from_clause - negated_from, _where, negated_params, _order, _score = await repo._build_fts_query_parts( - search_text="coffee NOT brewing", - ) - assert "AS fts_candidate" in negated_from - assert "FROM search_index AS candidate_all" in negated_from - assert negated_params["text_candidate"] == "coffee | brewing" + negated = compile_fts_filter(repo.scope, search_text="coffee NOT brewing") + assert "AS fts_candidate" in negated.from_clause + assert "FROM search_index AS candidate_all" in negated.from_clause + assert negated.params["text_candidate"] == "coffee | brewing" now = datetime.now(timezone.utc) rows = [ @@ -374,7 +376,11 @@ async def test_postgres_search_repository_tsquery_syntax_error_returns_empty( # Isolate database-error handling from the user parser, which deliberately # repairs malformed trailing operators before they reach PostgreSQL. with monkeypatch.context() as syntax_error: - syntax_error.setattr(repo, "_prepare_search_term", lambda *_args, **_kwargs: "coffee &") + syntax_error.setattr( + postgres_search_query_module, + "prepare_search_term", + lambda *_args, **_kwargs: "coffee &", + ) results = await repo.search(search_text="coffee") assert results == [] assert await repo.count(search_text="coffee") == 0 @@ -420,8 +426,8 @@ async def test_postgres_search_tsquery_error_does_not_poison_caller_session( # without the savepoint it aborts the caller's transaction. with monkeypatch.context() as syntax_error: syntax_error.setattr( - repo, - "_prepare_search_term", + postgres_search_query_module, + "prepare_search_term", lambda *_args, **_kwargs: "coffee &", ) results = await repo.search(search_text="coffee", session=session) @@ -1193,21 +1199,21 @@ async def test_postgres_question_punctuation_and_relaxation(session_maker, test_ and a strict all-AND miss had no relaxed retry, silently disabling the FTS half of hybrid search for natural-language questions. """ - repo = PostgresSearchRepository(session_maker, project_id=test_project.id) + PostgresSearchRepository(session_maker, project_id=test_project.id) # Edge punctuation stripped before lexeme formatting. - prepared = repo._prepare_search_term("When did Melanie paint a sunrise?") + prepared = prepare_search_term("When did Melanie paint a sunrise?") assert "?" not in prepared assert "sunrise:*" in prepared # Relaxation drops stopwords and OR-joins content terms. - relaxed = repo._relaxed_tsquery_text("When did Melanie paint a sunrise?") + relaxed = relaxed_tsquery_text("When did Melanie paint a sunrise?") assert relaxed == "melanie:* | paint:* | sunrise:*" # User intent is not second-guessed. - assert repo._relaxed_tsquery_text("alpha AND beta") is None - assert repo._relaxed_tsquery_text('"exact phrase"') is None - assert repo._relaxed_tsquery_text(None) is None + assert relaxed_tsquery_text("alpha AND beta") is None + assert relaxed_tsquery_text('"exact phrase"') is None + assert relaxed_tsquery_text(None) is None @pytest.mark.asyncio @@ -1267,7 +1273,7 @@ async def test_postgres_relaxes_after_strict_tsquery_syntax_error( ) syntax_errors: list[Exception] = [] - real_is_syntax_error = repo._is_tsquery_syntax_error + real_is_syntax_error = postgres_search_repository_module.is_tsquery_syntax_error def record_syntax_error(exc: Exception) -> bool: is_syntax_error = real_is_syntax_error(exc) @@ -1275,7 +1281,10 @@ def record_syntax_error(exc: Exception) -> bool: syntax_errors.append(exc) return is_syntax_error - monkeypatch.setattr(repo, "_is_tsquery_syntax_error", record_syntax_error) + # The repository module binds the classifier at import; patch it where it is read. + monkeypatch.setattr( + postgres_search_repository_module, "is_tsquery_syntax_error", record_syntax_error + ) query = "foo None: """Quoted user syntax must become a complete tsquery expression before SQL.""" - assert _make_repo()._prepare_search_term(query) == expected + assert prepare_search_term(query) == expected def test_postgres_many_quoted_groups_restore_atomically() -> None: @@ -632,4 +633,4 @@ def test_postgres_many_quoted_groups_restore_atomically() -> None: query = " OR ".join(f'"term{index} word{index}"' for index in range(11)) expected = " | ".join(f"(term{index} & word{index})" for index in range(11)) - assert _make_repo()._prepare_search_term(query) == expected + assert prepare_search_term(query) == expected diff --git a/tests/repository/test_search_file_path_prefix.py b/tests/repository/test_search_file_path_prefix.py index 885cd96ba..c6d2b3954 100644 --- a/tests/repository/test_search_file_path_prefix.py +++ b/tests/repository/test_search_file_path_prefix.py @@ -335,7 +335,7 @@ async def test_semantic_retrieval_honors_the_scope( def test_condition_is_one_shared_predicate_for_both_dialects(): """The SQL text and its parameters are backend-independent by construction. - Both `_build_fts_query_parts` implementations call this one helper, so the + Both `compile_fts_filter` implementations call this one helper, so the identical-behavior claim above is structural rather than a coincidence two hand-written predicates happen to share. """ diff --git a/tests/repository/test_search_relaxed_rendering.py b/tests/repository/test_search_relaxed_rendering.py index fc43d062c..93188ef9a 100644 --- a/tests/repository/test_search_relaxed_rendering.py +++ b/tests/repository/test_search_relaxed_rendering.py @@ -4,8 +4,8 @@ import pytest -from basic_memory.repository.postgres_search_repository import PostgresSearchRepository -from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository +from basic_memory.repository.postgres_search_query import relaxed_tsquery_text +from basic_memory.repository.sqlite_search_query import relaxed_fts_text CREATE_FTS = ( "CREATE VIRTUAL TABLE t USING fts5(" @@ -27,7 +27,7 @@ ) def test_sqlite_relaxed_text_quotes_only_terms_that_need_it(query: str, expected: str) -> None: """Apostrophe terms are quoted; every other term renders exactly as before.""" - assert SQLiteSearchRepository._relaxed_fts_text(query) == expected + assert relaxed_fts_text(query) == expected @pytest.mark.parametrize( @@ -47,7 +47,7 @@ def test_sqlite_relaxed_text_is_accepted_by_fts5(query: str) -> None: catches and turns into an empty result — the relaxed retry then silently contributes nothing, which is the failure this fallback exists to prevent. """ - relaxed = SQLiteSearchRepository._relaxed_fts_text(query) + relaxed = relaxed_fts_text(query) assert relaxed is not None connection = sqlite3.connect(":memory:") @@ -82,7 +82,7 @@ def test_sqlite_relaxed_text_bare_apostrophe_would_be_rejected() -> None: ) def test_postgres_relaxed_tsquery_quotes_apostrophe_lexemes(query: str, expected: str) -> None: """Postgres carries the same token shapes, so it needs the same escaping.""" - assert PostgresSearchRepository._relaxed_tsquery_text(query) == expected + assert relaxed_tsquery_text(query) == expected @pytest.mark.parametrize( @@ -105,7 +105,7 @@ def test_relaxed_terms_match_the_stored_note(document: str, query: str) -> None: try: connection.execute(CREATE_FTS) connection.execute("INSERT INTO t VALUES (?)", (document,)) - relaxed = SQLiteSearchRepository._relaxed_fts_text(query) + relaxed = relaxed_fts_text(query) assert relaxed is not None rows = connection.execute("SELECT rowid FROM t WHERE t MATCH ?", (relaxed,)).fetchall() finally: @@ -125,7 +125,7 @@ def test_relaxed_terms_match_either_stored_form(document: str) -> None: try: connection.execute(CREATE_FTS) connection.execute("INSERT INTO t VALUES (?)", (document,)) - relaxed = SQLiteSearchRepository._relaxed_fts_text("foo­bar права доступа") + relaxed = relaxed_fts_text("foo­bar права доступа") assert relaxed is not None rows = connection.execute("SELECT rowid FROM t WHERE t MATCH ?", (relaxed,)).fetchall() finally: @@ -138,5 +138,5 @@ def test_orthographic_joiners_are_not_duplicated_into_a_second_variant() -> None A stripped variant would only widen the OR with a term no note can hold. """ - relaxed = SQLiteSearchRepository._relaxed_fts_text("نمی‌خواهم دسترسی را لغو") + relaxed = relaxed_fts_text("نمی‌خواهم دسترسی را لغو") assert relaxed == "نمی‌خواهم* OR دسترسی* OR را* OR لغو*" diff --git a/tests/repository/test_search_repository.py b/tests/repository/test_search_repository.py index c73cd02c0..11b95f4e0 100644 --- a/tests/repository/test_search_repository.py +++ b/tests/repository/test_search_repository.py @@ -10,6 +10,7 @@ from basic_memory.models import Entity from basic_memory.models.project import Project from basic_memory.repository.search_repository import SearchIndexRow +from basic_memory.repository import postgres_search_query, sqlite_search_query from basic_memory.repository.postgres_search_repository import PostgresSearchRepository from basic_memory.schemas.search import SearchItemType @@ -19,6 +20,13 @@ def is_postgres_backend(search_repository): return isinstance(search_repository, PostgresSearchRepository) +def fts_query(search_repository): + """The term-preparation module for the repository's backend.""" + if is_postgres_backend(search_repository): + return postgres_search_query + return sqlite_search_query + + @pytest_asyncio.fixture async def search_entity(session_maker, test_project: Project): """Create a test entity for search testing.""" @@ -616,53 +624,57 @@ class TestSearchTermPreparation: def test_simple_terms_get_prefix_wildcard(self, search_repository): """Simple alphanumeric terms should get prefix matching.""" - from basic_memory.repository.postgres_search_repository import PostgresSearchRepository - - if isinstance(search_repository, PostgresSearchRepository): + if is_postgres_backend(search_repository): # Postgres tsquery uses :* for prefix matching - assert search_repository._prepare_search_term("hello") == "hello:*" - assert search_repository._prepare_search_term("project") == "project:*" - assert search_repository._prepare_search_term("test123") == "test123:*" + assert fts_query(search_repository).prepare_search_term("hello") == "hello:*" + assert fts_query(search_repository).prepare_search_term("project") == "project:*" + assert fts_query(search_repository).prepare_search_term("test123") == "test123:*" else: # SQLite FTS5 uses * for prefix matching - assert search_repository._prepare_search_term("hello") == "hello*" - assert search_repository._prepare_search_term("project") == "project*" - assert search_repository._prepare_search_term("test123") == "test123*" + assert fts_query(search_repository).prepare_search_term("hello") == "hello*" + assert fts_query(search_repository).prepare_search_term("project") == "project*" + assert fts_query(search_repository).prepare_search_term("test123") == "test123*" def test_terms_with_existing_wildcard_unchanged(self, search_repository): """Terms that already contain * should remain unchanged.""" if is_postgres_backend(search_repository): # Postgres uses different syntax (:* instead of *) - assert search_repository._prepare_search_term("hello*") == "hello:*" - assert search_repository._prepare_search_term("test*world") == "test:*world" + assert fts_query(search_repository).prepare_search_term("hello*") == "hello:*" + assert fts_query(search_repository).prepare_search_term("test*world") == "test:*world" else: - assert search_repository._prepare_search_term("hello*") == "hello*" - assert search_repository._prepare_search_term("test*world") == "test*world" + assert fts_query(search_repository).prepare_search_term("hello*") == "hello*" + assert fts_query(search_repository).prepare_search_term("test*world") == "test*world" def test_boolean_operators_preserved(self, search_repository): """Boolean operators should be preserved without modification.""" if is_postgres_backend(search_repository): # Postgres converts AND/OR/NOT to &/|/! - assert search_repository._prepare_search_term("hello AND world") == "hello & world" - assert search_repository._prepare_search_term("cat OR dog") == "cat | dog" + assert ( + fts_query(search_repository).prepare_search_term("hello AND world") + == "hello & world" + ) + assert fts_query(search_repository).prepare_search_term("cat OR dog") == "cat | dog" # NOT must be converted to "& !" for proper tsquery syntax assert ( - search_repository._prepare_search_term("project NOT meeting") + fts_query(search_repository).prepare_search_term("project NOT meeting") == "project & !meeting" ) assert ( - search_repository._prepare_search_term("(hello AND world) OR test") + fts_query(search_repository).prepare_search_term("(hello AND world) OR test") == "(hello & world) | test" ) else: - assert search_repository._prepare_search_term("hello AND world") == "hello AND world" - assert search_repository._prepare_search_term("cat OR dog") == "cat OR dog" assert ( - search_repository._prepare_search_term("project NOT meeting") + fts_query(search_repository).prepare_search_term("hello AND world") + == "hello AND world" + ) + assert fts_query(search_repository).prepare_search_term("cat OR dog") == "cat OR dog" + assert ( + fts_query(search_repository).prepare_search_term("project NOT meeting") == "project NOT meeting" ) assert ( - search_repository._prepare_search_term("(hello AND world) OR test") + fts_query(search_repository).prepare_search_term("(hello AND world) OR test") == "(hello AND world) OR test" ) @@ -672,30 +684,30 @@ def test_hyphenated_terms_with_boolean_operators(self, search_repository): pytest.skip("This test is for SQLite FTS5-specific quoting behavior") # Test the specific case from the GitHub issue - result = search_repository._prepare_search_term("tier1-test AND unicode") + result = fts_query(search_repository).prepare_search_term("tier1-test AND unicode") assert result == '"tier1-test" AND unicode' # Test other hyphenated Boolean combinations assert ( - search_repository._prepare_search_term("multi-word OR single") + fts_query(search_repository).prepare_search_term("multi-word OR single") == '"multi-word" OR single' ) assert ( - search_repository._prepare_search_term("well-formed NOT badly-formed") + fts_query(search_repository).prepare_search_term("well-formed NOT badly-formed") == '"well-formed" NOT "badly-formed"' ) assert ( - search_repository._prepare_search_term("test-case AND (hello OR world)") + fts_query(search_repository).prepare_search_term("test-case AND (hello OR world)") == '"test-case" AND (hello OR world)' ) # Test mixed special characters with Boolean operators assert ( - search_repository._prepare_search_term("config.json AND test-file") + fts_query(search_repository).prepare_search_term("config.json AND test-file") == '"config.json" AND "test-file"' ) assert ( - search_repository._prepare_search_term("C++ OR python-script") + fts_query(search_repository).prepare_search_term("C++ OR python-script") == '"C++" OR "python-script"' ) @@ -705,11 +717,14 @@ def test_programming_terms_should_work(self, search_repository): pytest.skip("This test is for SQLite FTS5-specific behavior") # These should be quoted to handle special characters safely - assert search_repository._prepare_search_term("C++") == '"C++"*' - assert search_repository._prepare_search_term("function()") == '"function()"*' - assert search_repository._prepare_search_term("email@domain.com") == '"email@domain.com"*' - assert search_repository._prepare_search_term("array[index]") == '"array[index]"*' - assert search_repository._prepare_search_term("config.json") == '"config.json"*' + assert fts_query(search_repository).prepare_search_term("C++") == '"C++"*' + assert fts_query(search_repository).prepare_search_term("function()") == '"function()"*' + assert ( + fts_query(search_repository).prepare_search_term("email@domain.com") + == '"email@domain.com"*' + ) + assert fts_query(search_repository).prepare_search_term("array[index]") == '"array[index]"*' + assert fts_query(search_repository).prepare_search_term("config.json") == '"config.json"*' def test_malformed_fts5_syntax_quoted(self, search_repository): """Malformed FTS5 syntax should be quoted to prevent errors.""" @@ -717,17 +732,21 @@ def test_malformed_fts5_syntax_quoted(self, search_repository): pytest.skip("This test is for SQLite FTS5-specific behavior") # Multiple operators without proper syntax - assert search_repository._prepare_search_term("+++invalid+++") == '"+++invalid+++"*' - assert search_repository._prepare_search_term("!!!error!!!") == '"!!!error!!!"*' - assert search_repository._prepare_search_term("@#$%^&*()") == '"@#$%^&*()"*' + assert ( + fts_query(search_repository).prepare_search_term("+++invalid+++") == '"+++invalid+++"*' + ) + assert fts_query(search_repository).prepare_search_term("!!!error!!!") == '"!!!error!!!"*' + assert fts_query(search_repository).prepare_search_term("@#$%^&*()") == '"@#$%^&*()"*' def test_quoted_strings_handled_properly(self, search_repository): """Strings with quotes should have quotes escaped.""" if is_postgres_backend(search_repository): pytest.skip("This test is for SQLite FTS5-specific behavior") - assert search_repository._prepare_search_term('say "hello"') == '"say ""hello"""*' - assert search_repository._prepare_search_term("it's working") == '"it\'s working"*' + assert fts_query(search_repository).prepare_search_term('say "hello"') == '"say ""hello"""*' + assert ( + fts_query(search_repository).prepare_search_term("it's working") == '"it\'s working"*' + ) def test_file_paths_no_prefix_wildcard(self, search_repository): """File paths should not get prefix wildcards.""" @@ -735,11 +754,11 @@ def test_file_paths_no_prefix_wildcard(self, search_repository): pytest.skip("This test is for SQLite FTS5-specific behavior") assert ( - search_repository._prepare_search_term("config.json", is_prefix=False) + fts_query(search_repository).prepare_search_term("config.json", is_prefix=False) == '"config.json"' ) assert ( - search_repository._prepare_search_term("docs/readme.md", is_prefix=False) + fts_query(search_repository).prepare_search_term("docs/readme.md", is_prefix=False) == '"docs/readme.md"' ) @@ -748,9 +767,12 @@ def test_spaces_handled_correctly(self, search_repository): if is_postgres_backend(search_repository): pytest.skip("This test is for SQLite FTS5-specific behavior") - assert search_repository._prepare_search_term("hello world") == "hello* AND world*" assert ( - search_repository._prepare_search_term("project planning") == "project* AND planning*" + fts_query(search_repository).prepare_search_term("hello world") == "hello* AND world*" + ) + assert ( + fts_query(search_repository).prepare_search_term("project planning") + == "project* AND planning*" ) def test_version_strings_with_dots_handled_correctly(self, search_repository): @@ -760,7 +782,7 @@ def test_version_strings_with_dots_handled_correctly(self, search_repository): # This reproduces the bug where "Basic Memory v0.13.0b2" becomes "Basic* AND Memory* AND v0.13.0b2*" # which causes FTS5 syntax errors because v0.13.0b2* is not valid FTS5 syntax - result = search_repository._prepare_search_term("Basic Memory v0.13.0b2") + result = fts_query(search_repository).prepare_search_term("Basic Memory v0.13.0b2") # Should be quoted because of dots in v0.13.0b2 assert result == '"Basic Memory v0.13.0b2"*' @@ -770,12 +792,18 @@ def test_mixed_special_characters_in_multi_word_queries(self, search_repository) pytest.skip("This test is for SQLite FTS5-specific behavior") # Any word containing special characters should cause the entire phrase to be quoted - assert search_repository._prepare_search_term("config.json file") == '"config.json file"*' assert ( - search_repository._prepare_search_term("user@email.com account") + fts_query(search_repository).prepare_search_term("config.json file") + == '"config.json file"*' + ) + assert ( + fts_query(search_repository).prepare_search_term("user@email.com account") == '"user@email.com account"*' ) - assert search_repository._prepare_search_term("node.js and react") == '"node.js and react"*' + assert ( + fts_query(search_repository).prepare_search_term("node.js and react") + == '"node.js and react"*' + ) @pytest.mark.asyncio async def test_search_with_special_characters_returns_results(self, search_repository): @@ -913,15 +941,19 @@ async def test_wildcard_only_search(self, search_repository, search_entity): def test_boolean_query_empty_parts_coverage(self, search_repository): """Test Boolean query parsing with empty parts (line 143 coverage).""" # Create queries that will result in empty parts after splitting - result1 = search_repository._prepare_boolean_query( + result1 = fts_query(search_repository).prepare_boolean_query( "hello AND AND world" ) # Double operator assert "hello" in result1 and "world" in result1 - result2 = search_repository._prepare_boolean_query(" OR test") # Leading operator + result2 = fts_query(search_repository).prepare_boolean_query( + " OR test" + ) # Leading operator assert "test" in result2 - result3 = search_repository._prepare_boolean_query("test OR ") # Trailing operator + result3 = fts_query(search_repository).prepare_boolean_query( + "test OR " + ) # Trailing operator assert "test" in result3 def test_parenthetical_term_quote_escaping(self, search_repository): @@ -930,12 +962,12 @@ def test_parenthetical_term_quote_escaping(self, search_repository): pytest.skip("This test is for SQLite FTS5-specific behavior") # Test term with quotes that needs escaping - result = search_repository._prepare_parenthetical_term('(say "hello" world)') + result = fts_query(search_repository).prepare_parenthetical_term('(say "hello" world)') # Should escape quotes by doubling them assert '""hello""' in result # Test term with single quotes - result2 = search_repository._prepare_parenthetical_term("(it's working)") + result2 = fts_query(search_repository).prepare_parenthetical_term("(it's working)") assert "it's working" in result2 def test_needs_quoting_empty_input(self, search_repository): @@ -944,26 +976,26 @@ def test_needs_quoting_empty_input(self, search_repository): pytest.skip("This test is for SQLite FTS5-specific behavior") # Test empty string - assert not search_repository._needs_quoting("") + assert not fts_query(search_repository).needs_quoting("") # Test whitespace-only string - assert not search_repository._needs_quoting(" ") + assert not fts_query(search_repository).needs_quoting(" ") # Test None-like cases - assert not search_repository._needs_quoting("\t") + assert not fts_query(search_repository).needs_quoting("\t") def test_prepare_single_term_empty_input(self, search_repository): """Test _prepare_single_term with empty inputs (line 227 coverage).""" # Test empty string - result1 = search_repository._prepare_single_term("") + result1 = fts_query(search_repository).prepare_single_term("") assert result1 == "" # Test whitespace-only string - result2 = search_repository._prepare_single_term(" ") + result2 = fts_query(search_repository).prepare_single_term(" ") assert result2 == " " # Should return as-is # Test string that becomes empty after strip - result3 = search_repository._prepare_single_term("\t\n") + result3 = fts_query(search_repository).prepare_single_term("\t\n") assert result3 == "\t\n" # Should return original @@ -1192,7 +1224,7 @@ async def test_question_punctuation_does_not_phrase_quote(search_repository): '"When did Melanie paint a sunrise?"*' — zero rows for any corpus — which silently disabled the FTS half of hybrid search for question queries. """ - prepared = search_repository._prepare_single_term("When did Melanie paint a sunrise?") + prepared = fts_query(search_repository).prepare_single_term("When did Melanie paint a sunrise?") assert '"' not in prepared # Prefix syntax differs by backend: FTS5 uses '*', tsquery uses ':*'. if is_postgres_backend(search_repository): @@ -1205,10 +1237,10 @@ async def test_question_punctuation_does_not_phrase_quote(search_repository): async def test_relaxed_query_drops_stopwords(search_repository): """Relaxation keys on content-bearing terms in each backend's syntax.""" if is_postgres_backend(search_repository): - relaxed = search_repository._relaxed_tsquery_text("When did Melanie paint a sunrise?") + relaxed = postgres_search_query.relaxed_tsquery_text("When did Melanie paint a sunrise?") assert relaxed == "melanie:* | paint:* | sunrise:*" else: - relaxed = search_repository._relaxed_fts_text("When did Melanie paint a sunrise?") + relaxed = sqlite_search_query.relaxed_fts_text("When did Melanie paint a sunrise?") assert relaxed == "melanie* OR paint* OR sunrise*" @@ -1216,14 +1248,14 @@ async def test_relaxed_query_drops_stopwords(search_repository): async def test_relaxed_query_preserves_punctuated_ascii_token_pieces(search_repository): """Hyphenated and slashed ASCII terms should relax using their regex token pieces.""" if is_postgres_backend(search_repository): - relaxed = search_repository._relaxed_tsquery_text("client-side state management") + relaxed = postgres_search_query.relaxed_tsquery_text("client-side state management") assert relaxed == "client:* | side:* | state:* | management:*" - slashed = search_repository._relaxed_tsquery_text("foo/bar baz qux") + slashed = postgres_search_query.relaxed_tsquery_text("foo/bar baz qux") assert slashed == "foo:* | bar:* | baz:* | qux:*" else: - relaxed = search_repository._relaxed_fts_text("client-side state management") + relaxed = sqlite_search_query.relaxed_fts_text("client-side state management") assert relaxed == "client* OR side* OR state* OR management*" - slashed = search_repository._relaxed_fts_text("foo/bar baz qux") + slashed = sqlite_search_query.relaxed_fts_text("foo/bar baz qux") assert slashed == "foo* OR bar* OR baz* OR qux*" @@ -1231,10 +1263,10 @@ async def test_relaxed_query_preserves_punctuated_ascii_token_pieces(search_repo async def test_relaxed_query_supports_whitespace_separated_cjk_terms(search_repository): """CJK terms separated by spaces should relax even when ASCII tokenization finds none.""" if is_postgres_backend(search_repository): - relaxed = search_repository._relaxed_tsquery_text("季度 报告") + relaxed = postgres_search_query.relaxed_tsquery_text("季度 报告") assert relaxed == "季度:* | 报告:*" else: - relaxed = search_repository._relaxed_fts_text("季度 报告") + relaxed = sqlite_search_query.relaxed_fts_text("季度 报告") assert relaxed == "季度* OR 报告*" @@ -1243,9 +1275,9 @@ async def test_relaxed_query_respects_user_intent(search_repository): # Eligibility matches the service-level relaxation (both backends): quoted, # boolean, short (<3 tokens), and numeric-identifier queries are not relaxed. relaxer = ( - search_repository._relaxed_tsquery_text + postgres_search_query.relaxed_tsquery_text if is_postgres_backend(search_repository) - else search_repository._relaxed_fts_text + else sqlite_search_query.relaxed_fts_text ) assert relaxer("alpha AND beta") is None assert relaxer('"exact phrase"') is None diff --git a/tests/repository/test_search_scope.py b/tests/repository/test_search_scope.py new file mode 100644 index 000000000..8120614a0 --- /dev/null +++ b/tests/repository/test_search_scope.py @@ -0,0 +1,47 @@ +"""ProjectScope: the explicit project set every search statement binds.""" + +from typing import Any, cast + +import pytest + +from basic_memory.repository.search_scope import ProjectScope + + +def test_of_sorts_and_dedupes() -> None: + assert ProjectScope.of([3, 1, 3, 2]).project_ids == (1, 2, 3) + assert ProjectScope.of([3, 1, 3, 2]) == ProjectScope.of((1, 2, 3)) + + +def test_single_and_empty() -> None: + assert ProjectScope.single(7).project_ids == (7,) + assert not ProjectScope.single(7).is_empty + assert ProjectScope.of([]).is_empty + + +@pytest.mark.parametrize("bad", [0, -1, True]) +def test_rejects_non_positive_ids(bad: int) -> None: + with pytest.raises(ValueError, match="positive integers"): + ProjectScope.of([bad]) + + +def test_rejects_non_int_ids() -> None: + with pytest.raises(ValueError, match="positive integers"): + ProjectScope.of(cast("list[int]", ["1"])) + + +def test_predicate_binds_each_id_once_per_statement() -> None: + params: dict[str, Any] = {} + scope = ProjectScope.of([5, 2]) + assert ( + scope.predicate("search_index.project_id", params) + == "search_index.project_id IN (:scope_0, :scope_1)" + ) + # A second reference within the same statement reuses the binds. + assert scope.predicate("owner.project_id", params) == "owner.project_id IN (:scope_0, :scope_1)" + assert params == {"scope_0": 2, "scope_1": 5} + + +def test_empty_scope_matches_nothing_and_binds_nothing() -> None: + params: dict[str, Any] = {} + assert ProjectScope.of([]).predicate("search_index.project_id", params) == "1 = 0" + assert params == {} diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 6e08fe31c..11a9030e1 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -80,10 +80,6 @@ async def record_entity_vector_deferrals( return None # no session_maker in this double; the real write is covered # in tests/services/test_project_readiness.py - @override - def _prepare_search_term(self, term, is_prefix=True): - return term - @override async def search( self, diff --git a/tests/repository/test_semantic_vector_sync.py b/tests/repository/test_semantic_vector_sync.py index 307b53030..a1669e464 100644 --- a/tests/repository/test_semantic_vector_sync.py +++ b/tests/repository/test_semantic_vector_sync.py @@ -45,10 +45,6 @@ async def init_search_index(self): async def get_entity_physical_chunk_keys(self, entity_id: int) -> set[str] | None: return None # physical storage is not inspectable in this double - @override - def _prepare_search_term(self, term, is_prefix=True): - return term - @override async def search( self, diff --git a/tests/repository/test_vector_pagination.py b/tests/repository/test_vector_pagination.py index 9139f4859..a75a7ce99 100644 --- a/tests/repository/test_vector_pagination.py +++ b/tests/repository/test_vector_pagination.py @@ -55,10 +55,6 @@ async def init_search_index(self): async def get_entity_physical_chunk_keys(self, entity_id: int) -> set[str] | None: return None # physical storage is not inspectable in this double - @override - def _prepare_search_term(self, term, is_prefix=True): - return term # pragma: no cover - @override async def search( self, diff --git a/tests/repository/test_vector_threshold.py b/tests/repository/test_vector_threshold.py index 1c6323aac..04bcf2c58 100644 --- a/tests/repository/test_vector_threshold.py +++ b/tests/repository/test_vector_threshold.py @@ -57,10 +57,6 @@ async def init_search_index(self): async def get_entity_physical_chunk_keys(self, entity_id: int) -> set[str] | None: return None # physical storage is not inspectable in this double - @override - def _prepare_search_term(self, term, is_prefix=True): - return term # pragma: no cover - @override async def search( self, From 51f95bfc6050916a398e487cd68673084462d2d4 Mon Sep 17 00:00:00 2001 From: Paul Hernandez <60959+phernandez@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:55:33 -0500 Subject: [PATCH 2/8] refactor(core): move FTS execution behind FtsBackend and bind the read path to scope (#1563) Signed-off-by: phernandez Co-authored-by: Claude Fable 5.1 --- CHANGELOG.md | 10 + .../repository/postgres_search_query.py | 276 ++++++++++++--- .../repository/postgres_search_repository.py | 329 +----------------- src/basic_memory/repository/search_filters.py | 230 ++++++++++-- .../repository/search_index_row.py | 6 + src/basic_memory/repository/search_query.py | 30 ++ .../repository/search_repository_base.py | 311 ++++++----------- src/basic_memory/repository/search_trace.py | 28 +- .../repository/sqlite_search_query.py | 280 ++++++++++++--- .../repository/sqlite_search_repository.py | 326 +---------------- .../services/project_readiness.py | 25 +- src/basic_memory/services/search_service.py | 23 +- tests/repository/test_hybrid_fusion.py | 4 + .../test_postgres_search_quoted_queries.py | 8 +- .../test_postgres_search_repository.py | 17 +- .../test_search_file_path_prefix.py | 2 +- tests/repository/test_search_trace.py | 9 +- tests/repository/test_semantic_search_base.py | 4 + tests/repository/test_semantic_vector_sync.py | 4 + ...est_vector_filter_candidate_restriction.py | 2 +- tests/repository/test_vector_pagination.py | 4 + tests/repository/test_vector_threshold.py | 4 + 22 files changed, 892 insertions(+), 1040 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4262bb491..a523a2777 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -121,6 +121,16 @@ `(project_id, ...)` identity. Project repositories call the compilers with a scope of one; no query behavior changes. First step of the shared single/multi-project reader. +- **#1558**: Full-text execution leaves the project repositories. `SQLiteFts` and + `PostgresFts` run a compiled statement for any `ProjectScope` and own their engine's + failure semantics (FTS5 syntax errors answer empty, Postgres retries a malformed strict + tsquery relaxed inside a savepoint). `SearchRepositoryBase.search` and `count` are + concrete: shared vector/hybrid dispatch, then the engine's `FtsBackend`. The base read + path binds every statement to the repository's scope (manifest hydration, candidate row + fetch, readiness and drop classification). `PreparedSearchQuery` moves to the repository + layer with defaults, and the filter helpers both backends share move from the base into + `search_filters`. No query behavior changes. + ## v0.23.2 (2026-08-25) diff --git a/src/basic_memory/repository/postgres_search_query.py b/src/basic_memory/repository/postgres_search_query.py index ea30070b7..0fa2fc339 100644 --- a/src/basic_memory/repository/postgres_search_query.py +++ b/src/basic_memory/repository/postgres_search_query.py @@ -1,32 +1,40 @@ -"""PostgreSQL tsquery preparation: term syntax and filter compilation. +"""PostgreSQL tsquery preparation and execution. -Pure functions over a ``ProjectScope`` and the caller's filters. Nothing here opens -a session or owns an index. +Term preparation and filter compilation are pure functions over a ``ProjectScope`` and +a ``PreparedSearchQuery``. ``PostgresFts`` runs the compiled statement and owns +tsquery's failure semantics. Nothing here initializes or mutates an index. """ import json import re +import time from collections.abc import Sequence -from datetime import datetime from typing import Any +import logfire +from loguru import logger +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from basic_memory import db from basic_memory.repository.metadata_filters import parse_metadata_filters from basic_memory.repository.script_ngrams import analyze_script_query from basic_memory.repository.search_filters import ( AFTER_DATE_ORDER_BY, POSTGRES_FILTER_DIALECT, CompiledFilter, - shared_filter_conditions, -) -from basic_memory.repository.search_query import relaxation_word_tokens, relaxed_query_words -from basic_memory.repository.search_repository_base import ( - SearchIndexKey, metadata_contains_like_condition, metadata_filter_content_type_condition, + shared_filter_conditions, +) +from basic_memory.repository.search_index_row import SearchIndexKey, SearchIndexRow +from basic_memory.repository.search_query import ( + PreparedSearchQuery, + relaxation_word_tokens, + relaxed_query_words, ) from basic_memory.repository.search_scope import ProjectScope -from basic_memory.schemas.search import SearchItemType -from basic_memory.temporal import TemporalFilter +from basic_memory.repository.search_trace import SearchTraceCollector, build_fts_page_stage _TSQUERY_OPERAND_PATTERN = re.compile(r"'(?:''|[^'])*'(?::\*)?|[^\s&|!()]+") _TSQUERY_WORD_PATTERN = re.compile(r"[^\W_]+(?:'[^\W_]+)?", re.UNICODE) @@ -36,6 +44,24 @@ # tsquery special characters that must not reach the parser as text. _TSQUERY_SPECIAL_CHARS = ("&", "|", "!", "(", ")", ":") +# Every FTS statement returns these columns plus a score. +_RESULT_COLUMNS = """ + search_index.project_id, + search_index.id, + search_index.title, + search_index.permalink, + search_index.file_path, + search_index.type, + search_index.metadata, + search_index.from_id, + search_index.to_id, + search_index.relation_type, + search_index.entity_id, + search_index.content_snippet, + search_index.category, + search_index.created_at, + search_index.updated_at""" + # --- Term preparation --- @@ -436,18 +462,8 @@ def _script_candidate_from_clause(scope: ProjectScope, params: dict[str, Any]) - def compile_fts_filter( scope: ProjectScope, + query: PreparedSearchQuery, *, - search_text: str | None = None, - permalink: str | None = None, - permalink_match: str | None = None, - title: str | None = None, - note_types: Sequence[str] | None = None, - after_date: datetime | None = None, - search_item_types: Sequence[SearchItemType] | None = None, - categories: Sequence[str] | None = None, - metadata_filters: dict[str, Any] | None = None, - file_path_prefix: str | None = None, - temporal: TemporalFilter | None = None, allow_relaxed: bool = False, candidate_keys: Sequence[SearchIndexKey] | None = None, ) -> CompiledFilter: @@ -458,21 +474,12 @@ def compile_fts_filter( """ params: dict[str, Any] = {} conditions = shared_filter_conditions( - scope, - params, - dialect=POSTGRES_FILTER_DIALECT, - permalink=permalink, - file_path_prefix=file_path_prefix, - candidate_keys=candidate_keys, - search_item_types=search_item_types, - categories=categories, - note_types=note_types, - after_date=after_date, - temporal=temporal, + scope, params, dialect=POSTGRES_FILTER_DIALECT, query=query, candidate_keys=candidate_keys ) from_clause = "search_index" document_vector: str | None = None script_tsqueries: list[str] = [] + search_text = query.search_text # Wildcard-only and blank text add no text condition: every row matches. if search_text and search_text.strip() not in ("", "*"): @@ -525,15 +532,15 @@ def compile_fts_filter( for index in range(len(script_tsqueries)) ) - if title: - params["title_text"] = prepare_search_term(title.strip(), is_prefix=False) + if query.title: + params["title_text"] = prepare_search_term(query.title.strip(), is_prefix=False) conditions.append( "to_tsvector('english', search_index.title) @@ to_tsquery('english', :title_text)" ) - if permalink_match: - permalink_text = permalink_match.lower().strip() - if "*" in permalink_match: + if query.permalink_match: + permalink_text = query.permalink_match.lower().strip() + if "*" in query.permalink_match: # ``*`` becomes the LIKE wildcard. params["permalink"] = permalink_text.replace("*", "%") conditions.append("search_index.permalink LIKE :permalink") @@ -543,8 +550,8 @@ def compile_fts_filter( # Structured metadata filters use jsonb_extract_path_text() / jsonb_extract_path() # with parameterized path parts instead of #>> / #> with interpolated paths. - if metadata_filters: - parsed_filters = parse_metadata_filters(metadata_filters) + if query.metadata_filters: + parsed_filters = parse_metadata_filters(query.metadata_filters) from_clause = f"{from_clause} JOIN entity ON search_index.entity_id = entity.id" # Frontmatter filters answer for notes only; see # metadata_filter_content_type_condition for why every regular file would @@ -662,6 +669,193 @@ def compile_fts_filter( from_clause=from_clause, where_clause=" AND ".join(conditions), params=params, - order_by_clause=AFTER_DATE_ORDER_BY if after_date else "", + order_by_clause=AFTER_DATE_ORDER_BY if query.after_date else "", score_expression=score_expression, ) + + +# --- Execution --- + + +class PostgresFts: + """Run tsquery statements for any scope in one database.""" + + def __init__(self, session_maker: async_sessionmaker[AsyncSession]) -> None: + self._session_maker = session_maker + + async def search( + self, + scope: ProjectScope, + query: PreparedSearchQuery, + *, + limit: int, + offset: int, + allow_relaxed: bool = False, + session: AsyncSession | None = None, + candidate_keys: Sequence[SearchIndexKey] | None = None, + trace: SearchTraceCollector | None = None, + ) -> list[SearchIndexRow]: + """Run one tsquery page, retrying a malformed or empty strict query relaxed.""" + search_text = query.search_text + compiled = compile_fts_filter( + scope, query, allow_relaxed=allow_relaxed, candidate_keys=candidate_keys + ) + params = compiled.params + params["limit"] = limit + params["offset"] = offset + + sql = f""" + SELECT{_RESULT_COLUMNS}, + {compiled.score_expression} as score + FROM {compiled.from_clause} + WHERE {compiled.where_clause} + ORDER BY score DESC {compiled.order_by_clause}, search_index.id ASC + LIMIT :limit + OFFSET :offset + """ + + logger.trace(f"Search {sql} params: {params}") + fts_started_at = time.perf_counter() if trace is not None else None + + use_savepoint = session is not None or allow_relaxed + + async def execute_rows(active_session: AsyncSession, query_params: dict[str, Any]): + # PostgreSQL leaves a transaction unusable after invalid tsquery syntax. + # Scope retryable or caller-owned attempts to a savepoint so a relaxed + # retry, and any caller continuing to use its session, starts healthy. + if use_savepoint: + async with active_session.begin_nested(): + result = await active_session.execute(text(sql), query_params) + return result.fetchall() + result = await active_session.execute(text(sql), query_params) + return result.fetchall() + + async def run_search(active_session: AsyncSession): + relaxed = relaxed_tsquery_text(search_text) if allow_relaxed else None + strict_syntax_error = False + relaxed_fallback_used = False + try: + rows = await execute_rows(active_session, params) + except Exception as exc: + if not (is_tsquery_syntax_error(exc) and relaxed and params.get("text")): + raise + strict_syntax_error = True + rows = [] + + # Trigger: multi-word natural-language query matched nothing under the + # default all-terms-AND tsquery semantics, or its punctuation produced + # invalid strict tsquery syntax. + # Why: questions rarely have every word in one document; without + # relaxation the FTS half of hybrid search contributes zero candidates. + # The relaxed renderer also tokenizes punctuation safely. + # Outcome: one retry with OR-joined prefix lexemes; ts_rank still ranks + # multi-term matches first. + if relaxed and not rows and params.get("text"): + relaxed_fallback_used = True + retry_reason = "invalid syntax" if strict_syntax_error else "0 results" + logger.debug( + f"Strict Postgres FTS returned {retry_reason}; retrying relaxed FTS query " + f"strict='{search_text}' relaxed='{relaxed}'" + ) + with logfire.span( + "search.relaxed_fts_retry", + backend="postgres", + reason="syntax_error" if strict_syntax_error else "empty_result", + token_count=len(relaxed_query_words(search_text) or ()), + limit=limit, + offset=offset, + ): + rows = await execute_rows(active_session, {**params, "text": relaxed}) + return rows, relaxed_fallback_used + + try: + if session is not None: + rows, relaxed_fallback_used = await run_search(session) + else: + async with db.scoped_session(self._session_maker) as owned_session: + rows, relaxed_fallback_used = await run_search(owned_session) + except Exception as e: + if is_tsquery_syntax_error(e): + logger.warning(f"tsquery syntax error for search term: {search_text}, error: {e}") + if trace is not None: + trace.fts = build_fts_page_stage( + [], + relaxed_fallback_used=False, + fts_ms=( + (time.perf_counter() - fts_started_at) * 1000 + if fts_started_at is not None + else None + ), + ) + return [] + logger.error(f"Database error during search: {e}") + raise + + results = [SearchIndexRow.from_mapping(row._asdict()) for row in rows] + if trace is not None: + trace.fts = build_fts_page_stage( + [((row.type, row.id), row.score or 0.0) for row in results], + relaxed_fallback_used=relaxed_fallback_used, + fts_ms=( + (time.perf_counter() - fts_started_at) * 1000 + if fts_started_at is not None + else None + ), + ) + + logger.trace(f"Found {len(results)} search results") + for r in results: + logger.trace( + f"Search result: project_id: {r.project_id} type:{r.type} title: {r.title} permalink: {r.permalink} score: {r.score}" + ) + return results + + async def count( + self, + scope: ProjectScope, + query: PreparedSearchQuery, + *, + allow_relaxed: bool = False, + ) -> int: + """Count rows matching the tsquery, with the same relaxed retry as search.""" + search_text = query.search_text + compiled = compile_fts_filter(scope, query, allow_relaxed=allow_relaxed) + params = compiled.params + sql = f"SELECT COUNT(*) FROM {compiled.from_clause} WHERE {compiled.where_clause}" + logger.trace(f"Count {sql} params: {params}") + + async def execute_count(active_session: AsyncSession, query_params: dict[str, Any]) -> int: + if allow_relaxed: + async with active_session.begin_nested(): + result = await active_session.execute(text(sql), query_params) + return int(result.scalar_one()) + result = await active_session.execute(text(sql), query_params) + return int(result.scalar_one()) + + try: + async with db.scoped_session(self._session_maker) as session: + relaxed = relaxed_tsquery_text(search_text) if allow_relaxed else None + strict_syntax_error = False + try: + total = await execute_count(session, params) + except Exception as exc: + if not (is_tsquery_syntax_error(exc) and relaxed and params.get("text")): + raise + strict_syntax_error = True + total = 0 + + if relaxed and total == 0 and params.get("text"): + with logfire.span( + "search.count.relaxed_fts_retry", + backend="postgres", + reason="syntax_error" if strict_syntax_error else "empty_result", + token_count=len(relaxed_query_words(search_text) or ()), + ): + total = await execute_count(session, {**params, "text": relaxed}) + return total + except Exception as e: + if is_tsquery_syntax_error(e): + logger.warning(f"tsquery syntax error for search term: {search_text}, error: {e}") + return 0 + logger.error(f"Database error during search count: {e}") + raise diff --git a/src/basic_memory/repository/postgres_search_repository.py b/src/basic_memory/repository/postgres_search_repository.py index fc89e18da..0cb17fcb0 100644 --- a/src/basic_memory/repository/postgres_search_repository.py +++ b/src/basic_memory/repository/postgres_search_repository.py @@ -2,12 +2,9 @@ import asyncio import json -import time from collections.abc import Sequence -from datetime import datetime -from typing import Any, override, List, Optional +from typing import Any, override, List -import logfire from loguru import logger from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession @@ -20,23 +17,13 @@ from basic_memory.repository.rerank_provider import RerankProvider from basic_memory.repository.rerank_provider_factory import create_rerank_provider from basic_memory.repository.search_index_row import SearchIndexRow -from basic_memory.repository.search_query import relaxed_query_words from basic_memory.repository.script_ngrams import build_script_ngrams from basic_memory.repository.semantic_chunking import VectorChunkRecord from basic_memory.repository.search_repository_base import ( - SearchIndexKey, SearchRepositoryBase, VectorChunkState, ) -from basic_memory.repository.search_trace import ( - SearchTraceCollector, - build_fts_page_stage, -) -from basic_memory.repository.postgres_search_query import ( - compile_fts_filter, - is_tsquery_syntax_error, - relaxed_tsquery_text, -) +from basic_memory.repository.postgres_search_query import PostgresFts from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError from basic_memory.repository.semantic_vector_index import SemanticVectorIndex from basic_memory.repository.semantic_vector_sync import ( @@ -49,8 +36,6 @@ ) from basic_memory.repository.pgvector_index import PgVectorIndex from basic_memory.repository.postgres_fts_chunks import split_postgres_fts_chunks -from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode -from basic_memory.temporal import TemporalFilter def _strip_nul_from_row(row_data: dict[str, Any]) -> dict[str, Any]: @@ -89,6 +74,7 @@ def __init__( rerank_provider: RerankProvider | None = None, ): super().__init__(session_maker, project_id) + self._fts = PostgresFts(session_maker) self._app_config = app_config or ConfigManager().config self._semantic_enabled = self._app_config.semantic_search_enabled self._semantic_vector_k = self._app_config.semantic_vector_k @@ -386,22 +372,6 @@ async def _ensure_vector_tables(self) -> None: logger.debug(f"Postgres vector tables ready (dimensions={self._vector_dimensions})") self._vector_tables_initialized = True - @override - async def _run_vector_query( - self, - session: AsyncSession, - query_embedding: list[float], - candidate_limit: int, - *, - trace: SearchTraceCollector | None = None, - ) -> list[dict[str, Any]]: - return await super()._run_vector_query( - session, - query_embedding, - candidate_limit, - trace=trace, - ) - @override def _vector_prepare_window_size(self) -> int: """Use a bounded config-driven prepare window for Postgres vector sync.""" @@ -611,296 +581,3 @@ async def bulk_index_items(self, search_index_rows: List[SearchIndexRow]) -> Non await self._replace_fts_chunks(session, search_index_rows) logger.debug(f"Bulk indexed {len(search_index_rows)} rows") await session.commit() - - # ------------------------------------------------------------------ - # FTS search (Postgres-specific) - # ------------------------------------------------------------------ - - @override - async def search( - self, - search_text: Optional[str] = None, - permalink: Optional[str] = None, - permalink_match: Optional[str] = None, - title: Optional[str] = None, - note_types: Optional[List[str]] = None, - after_date: Optional[datetime] = None, - search_item_types: Optional[List[SearchItemType]] = None, - categories: Optional[List[str]] = None, - metadata_filters: Optional[dict[str, Any]] = None, - file_path_prefix: Optional[str] = None, - temporal: Optional[TemporalFilter] = None, - retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, - min_similarity: Optional[float] = None, - limit: int = 10, - offset: int = 0, - allow_relaxed: bool = False, - session: AsyncSession | None = None, - *, - candidate_keys: Sequence[SearchIndexKey] | None = None, - trace: SearchTraceCollector | None = None, - ) -> List[SearchIndexRow]: - """Search across all indexed content using PostgreSQL tsvector.""" - # --- Dispatch vector / hybrid modes (shared logic) --- - dispatched = await self._dispatch_retrieval_mode( - search_text=search_text, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - retrieval_mode=retrieval_mode, - min_similarity=min_similarity, - limit=limit, - offset=offset, - trace=trace, - ) - if dispatched is not None: - return dispatched - - # --- FTS mode (Postgres-specific) --- - compiled = compile_fts_filter( - self.scope, - search_text=search_text, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - allow_relaxed=allow_relaxed, - candidate_keys=candidate_keys, - ) - params = compiled.params - params["limit"] = limit - params["offset"] = offset - - sql = f""" - SELECT - search_index.project_id, - search_index.id, - search_index.title, - search_index.permalink, - search_index.file_path, - search_index.type, - search_index.metadata, - search_index.from_id, - search_index.to_id, - search_index.relation_type, - search_index.entity_id, - search_index.content_snippet, - search_index.category, - search_index.created_at, - search_index.updated_at, - {compiled.score_expression} as score - FROM {compiled.from_clause} - WHERE {compiled.where_clause} - ORDER BY score DESC {compiled.order_by_clause}, search_index.id ASC - LIMIT :limit - OFFSET :offset - """ - - logger.trace(f"Search {sql} params: {params}") - fts_started_at = time.perf_counter() if trace is not None else None - - use_savepoint = session is not None or allow_relaxed - - async def execute_rows(active_session: AsyncSession, query_params: dict[str, Any]): - # PostgreSQL leaves a transaction unusable after invalid tsquery syntax. - # Scope retryable or caller-owned attempts to a savepoint so a relaxed - # retry—and any caller continuing to use its session—starts healthy. - if use_savepoint: - async with active_session.begin_nested(): - result = await active_session.execute(text(sql), query_params) - return result.fetchall() - result = await active_session.execute(text(sql), query_params) - return result.fetchall() - - async def run_search(active_session: AsyncSession): - relaxed = relaxed_tsquery_text(search_text) if allow_relaxed else None - strict_syntax_error = False - relaxed_fallback_used = False - try: - rows = await execute_rows(active_session, params) - except Exception as exc: - if not (is_tsquery_syntax_error(exc) and relaxed and params.get("text")): - raise - strict_syntax_error = True - rows = [] - - # Trigger: multi-word natural-language query matched nothing - # under the default all-terms-AND tsquery semantics, or its punctuation - # produced invalid strict tsquery syntax. - # Why: questions rarely have every word in one document; - # without relaxation the FTS half of hybrid search contributes zero - # candidates. The relaxed renderer also tokenizes punctuation safely. - # Outcome: one retry with OR-joined prefix lexemes; ts_rank - # still ranks multi-term matches first. - if relaxed and not rows and params.get("text"): - relaxed_fallback_used = True - retry_reason = "invalid syntax" if strict_syntax_error else "0 results" - logger.debug( - f"Strict Postgres FTS returned {retry_reason}; retrying relaxed FTS query " - f"strict='{search_text}' relaxed='{relaxed}'" - ) - with logfire.span( - "search.relaxed_fts_retry", - backend="postgres", - reason="syntax_error" if strict_syntax_error else "empty_result", - token_count=len(relaxed_query_words(search_text) or ()), - limit=limit, - offset=offset, - ): - rows = await execute_rows( - active_session, - {**params, "text": relaxed}, - ) - return rows, relaxed_fallback_used - - try: - if session is not None: - rows, relaxed_fallback_used = await run_search(session) - else: - async with db.scoped_session(self.session_maker) as owned_session: - rows, relaxed_fallback_used = await run_search(owned_session) - except Exception as e: - if is_tsquery_syntax_error(e): - logger.warning(f"tsquery syntax error for search term: {search_text}, error: {e}") - if trace is not None: - trace.fts = build_fts_page_stage( - [], - relaxed_fallback_used=False, - fts_ms=( - (time.perf_counter() - fts_started_at) * 1000 - if fts_started_at is not None - else None - ), - ) - return [] - - # Re-raise other database errors - logger.error(f"Database error during search: {e}") - raise - - results = [SearchIndexRow.from_mapping(row._asdict()) for row in rows] - if trace is not None: - trace.fts = build_fts_page_stage( - [((row.type, row.id), row.score or 0.0) for row in results], - relaxed_fallback_used=relaxed_fallback_used, - fts_ms=( - (time.perf_counter() - fts_started_at) * 1000 - if fts_started_at is not None - else None - ), - ) - - logger.trace(f"Found {len(results)} search results") - for r in results: - logger.trace( - f"Search result: project_id: {r.project_id} type:{r.type} title: {r.title} permalink: {r.permalink} score: {r.score}" - ) - - return results - - @override - async def count( - self, - search_text: Optional[str] = None, - permalink: Optional[str] = None, - permalink_match: Optional[str] = None, - title: Optional[str] = None, - note_types: Optional[List[str]] = None, - after_date: Optional[datetime] = None, - search_item_types: Optional[List[SearchItemType]] = None, - categories: Optional[List[str]] = None, - metadata_filters: Optional[dict[str, Any]] = None, - file_path_prefix: Optional[str] = None, - temporal: Optional[TemporalFilter] = None, - retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, - min_similarity: Optional[float] = None, - allow_relaxed: bool = False, - ) -> int: - """Count indexed content matching the Postgres FTS query.""" - if retrieval_mode != SearchRetrievalMode.FTS: - return await super().count( - search_text=search_text, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - retrieval_mode=retrieval_mode, - min_similarity=min_similarity, - ) - - compiled = compile_fts_filter( - self.scope, - search_text=search_text, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - allow_relaxed=allow_relaxed, - ) - params = compiled.params - sql = f"SELECT COUNT(*) FROM {compiled.from_clause} WHERE {compiled.where_clause}" - logger.trace(f"Count {sql} params: {params}") - - async def execute_count(active_session: AsyncSession, query_params: dict[str, Any]) -> int: - if allow_relaxed: - async with active_session.begin_nested(): - result = await active_session.execute(text(sql), query_params) - return int(result.scalar_one()) - result = await active_session.execute(text(sql), query_params) - return int(result.scalar_one()) - - try: - async with db.scoped_session(self.session_maker) as session: - relaxed = relaxed_tsquery_text(search_text) if allow_relaxed else None - strict_syntax_error = False - try: - total = await execute_count(session, params) - except Exception as exc: - if not (is_tsquery_syntax_error(exc) and relaxed and params.get("text")): - raise - strict_syntax_error = True - total = 0 - - if relaxed and total == 0 and params.get("text"): - with logfire.span( - "search.count.relaxed_fts_retry", - backend="postgres", - reason="syntax_error" if strict_syntax_error else "empty_result", - token_count=len(relaxed_query_words(search_text) or ()), - ): - total = await execute_count( - session, - {**params, "text": relaxed}, - ) - return total - except Exception as e: - if is_tsquery_syntax_error(e): - logger.warning(f"tsquery syntax error for search term: {search_text}, error: {e}") - return 0 - logger.error(f"Database error during search count: {e}") - raise diff --git a/src/basic_memory/repository/search_filters.py b/src/basic_memory/repository/search_filters.py index 0a3233bca..2431d11cb 100644 --- a/src/basic_memory/repository/search_filters.py +++ b/src/basic_memory/repository/search_filters.py @@ -1,35 +1,40 @@ -"""WHERE-clause pieces both search backends share. +"""WHERE-clause pieces both search backends share, and the FTS execution contract. The two FTS engines differ in how they match text and read JSON. Every other filter a search accepts asks the same question of the same columns on both, so it is compiled once here. A backend supplies the two spellings that differ through ``FilterDialect`` and appends its own text, title, permalink-pattern, and metadata predicates around the -shared ones. +shared ones. ``FtsBackend`` is the narrow contract through which the shared read path +runs a compiled full-text statement on one engine. """ from collections.abc import Sequence from dataclasses import dataclass -from datetime import datetime -from typing import Any +from typing import Any, Protocol + +from sqlalchemy.ext.asyncio import AsyncSession from basic_memory.repository.note_type_filters import ( POSTGRES_NOTE_TYPE_VALUE, SQLITE_NOTE_TYPE_VALUE, build_note_type_predicate, ) -from basic_memory.repository.search_repository_base import ( - SearchIndexKey, - candidate_key_restriction_condition, - file_path_prefix_condition, -) +from basic_memory.repository.search_index_row import SearchIndexKey, SearchIndexRow +from basic_memory.repository.search_query import PreparedSearchQuery from basic_memory.repository.search_scope import ProjectScope +from basic_memory.repository.search_trace import SearchTraceCollector from basic_memory.repository.temporal_filters import build_temporal_predicate -from basic_memory.schemas.search import SearchItemType -from basic_memory.temporal import TemporalFilter +from basic_memory.runtime.storage import RUNTIME_MARKDOWN_CONTENT_TYPE +from basic_memory.schemas.search import normalize_file_path_prefix # Newest edits first whenever the caller filtered on ``after_date``. AFTER_DATE_ORDER_BY = ", search_index.updated_at DESC" +# SQLite's LIKE has no default escape character, and Postgres's is already the +# backslash, so naming this one explicitly in every pattern is what lets a single +# escaped pattern mean the same thing on both backends. +_LIKE_ESCAPE_CHARACTER = "\\" + @dataclass(frozen=True, slots=True) class FilterDialect: @@ -65,19 +70,179 @@ class CompiledFilter: score_expression: str +class FtsBackend(Protocol): + """Run one engine's full-text statement for a scope and a prepared query. + + Both implementations compile through their own ``compile_fts_filter``, execute, and + own the engine's failure semantics: FTS5 syntax errors answer with no rows, while + Postgres retries a malformed strict tsquery through the relaxed renderer inside a + savepoint so the caller's transaction survives. + """ + + async def search( + self, + scope: ProjectScope, + query: PreparedSearchQuery, + *, + limit: int, + offset: int, + allow_relaxed: bool = False, + session: AsyncSession | None = None, + candidate_keys: Sequence[SearchIndexKey] | None = None, + trace: SearchTraceCollector | None = None, + ) -> list[SearchIndexRow]: ... + + async def count( + self, + scope: ProjectScope, + query: PreparedSearchQuery, + *, + allow_relaxed: bool = False, + ) -> int: ... + + +# --- Filters both backends compile identically --- + + +def file_path_prefix_condition( + file_path_prefix: str | None, + params: dict[str, Any], +) -> str | None: + """Build the SQL scoping search rows to one directory subtree of the project. + + One implementation, shared verbatim by both backends: a subtree scope that + means different things on SQLite and Postgres would report an exact total + for a match set the other dialect never produces. + + Boundary: the compared prefix carries its trailing separator, so "specs" + admits "specs/api.md" and never "specs-archive/api.md". + + Why an explicit-length comparison rather than ``file_path LIKE 'specs/%'``: + LIKE reads "_" and "%" as wildcards and both are ordinary characters in a + directory name, so "my_notes" would silently also admit "my-notes"; and + LIKE case-folds differently per backend, so one filter would answer two + different questions. SUBSTR equality has no pattern language to escape and + compares under each backend's deterministic default text collation, which is + byte equality on both, so the dialects match exactly the same rows. + """ + normalized = normalize_file_path_prefix(file_path_prefix) + if normalized is None: + return None + prefix = f"{normalized}/" + params["file_path_prefix"] = prefix + params["file_path_prefix_length"] = len(prefix) + return "SUBSTR(search_index.file_path, 1, :file_path_prefix_length) = :file_path_prefix" + + +def metadata_filter_content_type_condition(params: dict[str, Any]) -> str: + """Build the SQL restricting a metadata-filtered query to Markdown notes. + + Frontmatter is a Markdown-only construct, but every indexed file (PDF, image, + binary) gets its own ENTITY row whose ``entity_metadata`` carries no keys at all. + A positive predicate can never match one, so this constraint was invisible until + ``{"key": None}`` arrived: ``IS NULL`` is satisfied by the *absence* of a key, + which is exactly the state every regular file is in, and the whole non-note half + of a project counted into an exact total. + + Applied to any metadata filter, not just the null one, so the frontmatter-only + contract is a property of the clause rather than of which operator happened to + be used. + """ + params["metadata_filter_content_type"] = RUNTIME_MARKDOWN_CONTENT_TYPE + return "entity.content_type = :metadata_filter_content_type" + + +def metadata_contains_like_condition( + extract_expr: str, + value: Any, + *, + param_prefix: str, + params: dict[str, Any], +) -> str: + """Build the compatibility half of an array-contains metadata filter. + + The primary half of a ``{"tags": ["security"]}`` filter asks JSON whether the + array holds the element (``json_each`` on SQLite, ``@>`` on Postgres) and answers + only when the stored value really is a JSON array. Frontmatter written before + tags were normalized can hold the array's *text* instead, either JSON-quoted + ('["security", "auth"]') or as a Python repr ("['security', 'auth']"), and only a + substring match finds an element inside those. Hence a pattern per quote style. + + LIKE reads "%" and "_" in the searched-for value as wildcards, so interpolating + the value raw turned `tags has 100%` into a pattern that also matched + "100-percent". Escaping both wildcards and the escape character itself makes the + value literal again. + """ + escaped = ( + str(value) + .replace(_LIKE_ESCAPE_CHARACTER, _LIKE_ESCAPE_CHARACTER * 2) + .replace("%", f"{_LIKE_ESCAPE_CHARACTER}%") + .replace("_", f"{_LIKE_ESCAPE_CHARACTER}_") + ) + double_quoted_param = f"{param_prefix}_like" + single_quoted_param = f"{param_prefix}_like_single" + params[double_quoted_param] = f'%"{escaped}"%' + params[single_quoted_param] = f"%'{escaped}'%" + escape_clause = f" ESCAPE '{_LIKE_ESCAPE_CHARACTER}'" + return ( + f"{extract_expr} LIKE :{double_quoted_param}{escape_clause} " + f"OR {extract_expr} LIKE :{single_quoted_param}{escape_clause}" + ) + + +def candidate_key_restriction_condition( + candidate_keys: Sequence[SearchIndexKey], + params: dict[str, Any], +) -> str: + """Build the SQL restricting a filter query to an explicit set of search rows. + + This is what turns the vector/hybrid filter pass from "give me a page of everything + the filter admits" into "of *these* candidates, which does the filter admit". The + first question has an answer the size of the project and had to be capped, and every + candidate outside the cap was then read as disallowed (#1431). The second question's + answer is bounded by the candidate set itself, so no cap is needed and none of the + candidates can fall off the end. + + Keys are grouped by row type rather than emitted as one ``(type, id)`` pair per + branch: entity, observation, and relation ids come from independent sequences, so the + type is part of the identity, but a handful of type-scoped ``IN`` lists binds one + parameter per key instead of two and leaves the id list in the shape both planners + can drive an index from. PostgreSQL's ``search_index`` primary key is + ``(id, type, project_id)``. + + An empty candidate set is a real state, not a caller error (a vector search whose + every hit was already dropped), and it admits nothing, so it yields a false + predicate rather than the vacuous truth an empty ``OR`` would collapse to. + """ + ids_by_type: dict[str, list[int]] = {} + for row_type, row_id in candidate_keys: + ids_by_type.setdefault(row_type, []).append(row_id) + + branches: list[str] = [] + for type_index, (row_type, row_ids) in enumerate(ids_by_type.items()): + type_param = f"candidate_type_{type_index}" + params[type_param] = row_type + id_params: list[str] = [] + for id_index, row_id in enumerate(dict.fromkeys(row_ids)): + id_param = f"candidate_id_{type_index}_{id_index}" + params[id_param] = row_id + id_params.append(f":{id_param}") + branches.append( + f"(search_index.type = :{type_param} AND search_index.id IN ({', '.join(id_params)}))" + ) + + if not branches: + return "1 = 0" + return f"({' OR '.join(branches)})" + + def shared_filter_conditions( scope: ProjectScope, params: dict[str, Any], *, dialect: FilterDialect, - permalink: str | None, - file_path_prefix: str | None, + query: PreparedSearchQuery, candidate_keys: Sequence[SearchIndexKey] | None, - search_item_types: Sequence[SearchItemType] | None, - categories: Sequence[str] | None, - note_types: Sequence[str] | None, - after_date: datetime | None, - temporal: TemporalFilter | None, ) -> list[str]: """Compile the filters whose SQL is identical on both backends. @@ -87,23 +252,20 @@ def shared_filter_conditions( """ conditions = [scope.predicate("search_index.project_id", params)] - if permalink: - params["permalink"] = permalink + if query.permalink: + params["permalink"] = query.permalink conditions.append("search_index.permalink = :permalink") - # See file_path_prefix_condition for the subtree boundary and escaping rules. - subtree_condition = file_path_prefix_condition(file_path_prefix, params) + subtree_condition = file_path_prefix_condition(query.file_path_prefix, params) if subtree_condition is not None: conditions.append(subtree_condition) - # See candidate_key_restriction_condition for why the vector filter pass asks - # about its candidates rather than paging the filter's whole match set (#1431). if candidate_keys is not None: conditions.append(candidate_key_restriction_condition(candidate_keys, params)) - if search_item_types: + if query.search_item_types: type_placeholders: list[str] = [] - for index, item_type in enumerate(search_item_types): + for index, item_type in enumerate(query.search_item_types): name = f"search_type_{index}" params[name] = item_type.value type_placeholders.append(f":{name}") @@ -114,9 +276,9 @@ def shared_filter_conditions( # callers expect exact-category matching, not incidental text matches. # Outcome: only rows whose indexed category exactly equals a requested value # survive (entities/relations have NULL category and are excluded). - if categories: + if query.categories: category_placeholders: list[str] = [] - for index, category in enumerate(categories): + for index, category in enumerate(query.categories): name = f"category_{index}" params[name] = category category_placeholders.append(f":{name}") @@ -125,23 +287,23 @@ def shared_filter_conditions( # The note type belongs to the note, but only its entity row carries the # frontmatter, so the predicate resolves through the owning note. See # note_type_filters for why reading it off each row excluded every non-entity row. - if note_types: + if query.note_types: conditions.append( build_note_type_predicate( - note_types, params, scope=scope, note_type_value=dialect.note_type_value + query.note_types, params, scope=scope, note_type_value=dialect.note_type_value ) ) # Filter on updated_at so recently edited notes are included even when created_at # is old. The matching ORDER BY lives in AFTER_DATE_ORDER_BY. - if after_date: - params["after_date"] = after_date + if query.after_date: + params["after_date"] = query.after_date conditions.append(dialect.after_date_condition) # Authored valid time (SPEC-82) is independent of ``after_date``: that one is # bookkeeping about the file, this one is a claim about the world. See # temporal_filters for the overlap rule and why the subquery is non-correlated. - if temporal is not None: - conditions.append(build_temporal_predicate(temporal, params, scope=scope)) + if query.temporal is not None: + conditions.append(build_temporal_predicate(query.temporal, params, scope=scope)) return conditions diff --git a/src/basic_memory/repository/search_index_row.py b/src/basic_memory/repository/search_index_row.py index ce8dc8d55..e9835b6e0 100644 --- a/src/basic_memory/repository/search_index_row.py +++ b/src/basic_memory/repository/search_index_row.py @@ -156,3 +156,9 @@ def to_insert(self, serialize_json: bool = True): "updated_at": self.updated_at if self.updated_at else None, "project_id": self.project_id, } + + +# Entity, observation, and relation rows carry ids from independent auto-increment +# sequences, so a bare id is ambiguous across row types. Every map in the retrieval +# path keys rows by (type, id) to avoid collisions. +type SearchIndexKey = tuple[str, int] diff --git a/src/basic_memory/repository/search_query.py b/src/basic_memory/repository/search_query.py index d73d70c0c..9cd2cb2b4 100644 --- a/src/basic_memory/repository/search_query.py +++ b/src/basic_memory/repository/search_query.py @@ -2,6 +2,36 @@ import re import unicodedata +from dataclasses import dataclass +from datetime import datetime +from typing import Any + +from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode +from basic_memory.temporal import TemporalFilter + + +@dataclass(frozen=True) +class PreparedSearchQuery: + """Normalized query inputs shared by search and count. + + Built once at the service boundary from the API's ``SearchQuery``; every layer + below reads the same value instead of threading thirteen keyword arguments. + """ + + search_text: str | None = None + permalink: str | None = None + permalink_match: str | None = None + title: str | None = None + note_types: list[str] | None = None + search_item_types: list[SearchItemType] | None = None + categories: list[str] | None = None + after_date: datetime | None = None + metadata_filters: dict[str, Any] | None = None + file_path_prefix: str | None = None + temporal: TemporalFilter | None = None + retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS + min_similarity: float | None = None + # Interrogative/function words contribute lexical noise when a strict # full-text query is relaxed: "when OR did OR a" matches loud wrong documents diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 99f8b23ae..31f97096a 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -32,7 +32,9 @@ demote_tail_scores, validate_rerank_scores, ) -from basic_memory.repository.search_index_row import SearchIndexRow +from basic_memory.repository.search_filters import FtsBackend +from basic_memory.repository.search_index_row import SearchIndexKey, SearchIndexRow +from basic_memory.repository.search_query import PreparedSearchQuery from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.script_ngrams import build_script_ngrams from basic_memory.repository.search_trace import ( @@ -78,12 +80,10 @@ StagedVectorDeletion as _StagedVectorDeletion, VectorChunkState, ) -from basic_memory.runtime.storage import RUNTIME_MARKDOWN_CONTENT_TYPE from basic_memory.runtime.vector_sync import VectorSyncBatchResult from basic_memory.schemas.search import ( SearchItemType, SearchRetrievalMode, - normalize_file_path_prefix, ) from basic_memory.temporal import TemporalFilter from basic_memory.utils import ensure_timezone_aware @@ -102,13 +102,19 @@ # vector-index change, or one still pending. Readiness reporting must apply the same # predicate — calling such a row "embedded" would report an index settled that # retrieval cannot answer from, which is the class of lie #1414 exists to remove. -# Both callers bind :project_id, :vector_index, and :embedding_model. -CURRENT_VECTOR_MANIFEST_PREDICATE = ( - "project_id = :project_id " - "AND vector_index = :vector_index " - "AND embedding_model = :embedding_model " - "AND embedding_status = 'ready'" -) +# Callers bind :vector_index and :embedding_model; the scope binds its own IDs. + + +def current_vector_manifest_predicate(scope: ProjectScope, params: dict[str, Any]) -> str: + """SQL admitting only manifest rows retrieval can answer from, within ``scope``.""" + return ( + f"{scope.predicate('project_id', params)} " + "AND vector_index = :vector_index " + "AND embedding_model = :embedding_model " + "AND embedding_status = 'ready'" + ) + + # Over-fetch factor for the rerank candidate chunk pool: chunks collapse to unique # (type, id) rows before reranking, so fetch several times reranker_candidates chunks # to keep enough unique documents in the rerank window. @@ -122,10 +128,6 @@ _SQLITE_MAX_PREPARE_WINDOW = semantic_vector_sync.SQLITE_MAX_PREPARE_WINDOW _BUILT_IN_VECTOR_INDEX_NAMES = frozenset({"pgvector", "sqlite-vec"}) -# Entity, observation, and relation rows in search_index carry ids from independent -# auto-increment sequences, so a bare id is ambiguous across row types. Every map in -# the vector/hybrid retrieval path must key rows by (type, id) to avoid collisions. -type SearchIndexKey = tuple[str, int] type StoredEmbeddingStatus = Literal["pending", "ready"] @@ -175,158 +177,6 @@ def __post_init__(self) -> None: object.__setattr__(self, "updated_at", ensure_timezone_aware(updated_at)) -def file_path_prefix_condition( - file_path_prefix: Optional[str], - params: Dict[str, Any], -) -> Optional[str]: - """Build the SQL scoping search rows to one directory subtree of the project. - - One implementation, shared verbatim by both backends: a subtree scope that - means different things on SQLite and Postgres would report an exact total - for a match set the other dialect never produces. - - Boundary: the compared prefix carries its trailing separator, so "specs" - admits "specs/api.md" and never "specs-archive/api.md". - - Why an explicit-length comparison rather than ``file_path LIKE 'specs/%'``: - LIKE reads "_" and "%" as wildcards and both are ordinary characters in a - directory name, so "my_notes" would silently also admit "my-notes"; and - LIKE case-folds differently per backend — SQLite's is ASCII-case-insensitive - while Postgres's is case-sensitive — so one filter would answer two - different questions. SUBSTR equality has no pattern language to escape and - compares under each backend's deterministic default text collation, which is - byte equality on both, so the dialects match exactly the same rows. - """ - normalized = normalize_file_path_prefix(file_path_prefix) - if normalized is None: - return None - prefix = f"{normalized}/" - params["file_path_prefix"] = prefix - params["file_path_prefix_length"] = len(prefix) - return "SUBSTR(search_index.file_path, 1, :file_path_prefix_length) = :file_path_prefix" - - -def metadata_filter_content_type_condition(params: Dict[str, Any]) -> str: - """Build the SQL restricting a metadata-filtered query to Markdown notes. - - Frontmatter is a Markdown-only construct, but every indexed file — PDF, - image, binary — gets its own ENTITY row whose ``entity_metadata`` carries no - keys at all. A positive predicate can never match one, so this constraint - was invisible until ``{"key": None}`` arrived: ``IS NULL`` is satisfied by - the *absence* of a key, which is exactly the state every regular file is in, - and the whole non-note half of a project counted into an exact total. - - Applied to any metadata filter, not just the null one, so the - frontmatter-only contract is a property of the clause rather than of which - operator happened to be used. Shared by both backends for the same reason - the subtree scope is: a filter that admits different rows per dialect would - report an exact total for a match set the other never produces. - """ - params["metadata_filter_content_type"] = RUNTIME_MARKDOWN_CONTENT_TYPE - return "entity.content_type = :metadata_filter_content_type" - - -# SQLite's LIKE has no default escape character, and Postgres's is already the -# backslash, so naming this one explicitly in every pattern is what lets a single -# escaped pattern mean the same thing on both backends. -_LIKE_ESCAPE_CHARACTER = "\\" - - -def metadata_contains_like_condition( - extract_expr: str, - value: Any, - *, - param_prefix: str, - params: Dict[str, Any], -) -> str: - """Build the compatibility half of an array-contains metadata filter. - - The primary half of a ``{"tags": ["security"]}`` filter asks JSON whether the - array holds the element — ``json_each`` on SQLite, ``@>`` on Postgres — and - answers only when the stored value really is a JSON array. Frontmatter - written before tags were normalized can hold the array's *text* instead, - either JSON-quoted ('["security", "auth"]') or as a Python repr - ("['security', 'auth']"), and only a substring match finds an element inside - those. Hence a pattern per quote style, and hence the pattern-language - problem this function exists to solve. - - LIKE reads "%" and "_" in the searched-for value as wildcards, so - interpolating the value raw turned `tags has 100%` into a pattern that also - matched "100-percent" — a wrong hit and an inflated exact total, produced by - the branch the caller only meant as a fallback. Escaping both wildcards and - the escape character itself makes the value literal again. - - Shared by both backends for the same reason the subtree scope is: a filter - that admits different rows per dialect would report an exact total for a - match set the other never produces. - """ - escaped = ( - str(value) - .replace(_LIKE_ESCAPE_CHARACTER, _LIKE_ESCAPE_CHARACTER * 2) - .replace("%", f"{_LIKE_ESCAPE_CHARACTER}%") - .replace("_", f"{_LIKE_ESCAPE_CHARACTER}_") - ) - double_quoted_param = f"{param_prefix}_like" - single_quoted_param = f"{param_prefix}_like_single" - params[double_quoted_param] = f'%"{escaped}"%' - params[single_quoted_param] = f"%'{escaped}'%" - escape_clause = f" ESCAPE '{_LIKE_ESCAPE_CHARACTER}'" - return ( - f"{extract_expr} LIKE :{double_quoted_param}{escape_clause} " - f"OR {extract_expr} LIKE :{single_quoted_param}{escape_clause}" - ) - - -def candidate_key_restriction_condition( - candidate_keys: Sequence[SearchIndexKey], - params: Dict[str, Any], -) -> str: - """Build the SQL restricting a filter query to an explicit set of search rows. - - This is what turns the vector/hybrid filter pass from "give me a page of everything - the filter admits" into "of *these* candidates, which does the filter admit". The - first question has an answer the size of the project and had to be capped, and every - candidate outside the cap was then read as disallowed (#1431). The second question's - answer is bounded by the candidate set itself, so no cap is needed and none of the - candidates can fall off the end. - - Keys are grouped by row type rather than emitted as one ``(type, id)`` pair per - branch: entity, observation, and relation ids come from independent sequences, so the - type is part of the identity, but a handful of type-scoped ``IN`` lists binds one - parameter per key instead of two and leaves the id list in the shape both planners - can drive an index from. PostgreSQL's ``search_index`` primary key is - ``(id, type, project_id)``. - - An empty candidate set is a real state, not a caller error — a vector search whose - every hit was already dropped — and it admits nothing, so it yields a false - predicate rather than the vacuous truth an empty ``OR`` would collapse to. - - Shared verbatim by both backends for the same reason the subtree scope is: a - restriction that admitted different rows per dialect would give semantic search a - different candidate set depending on which database happened to be underneath. - """ - ids_by_type: dict[str, list[int]] = {} - for row_type, row_id in candidate_keys: - ids_by_type.setdefault(row_type, []).append(row_id) - - branches: list[str] = [] - for type_index, (row_type, row_ids) in enumerate(ids_by_type.items()): - type_param = f"candidate_type_{type_index}" - params[type_param] = row_type - id_params: list[str] = [] - for id_index, row_id in enumerate(dict.fromkeys(row_ids)): - id_param = f"candidate_id_{type_index}_{id_index}" - params[id_param] = row_id - id_params.append(f":{id_param}") - branches.append( - f"(search_index.type = :{type_param} AND search_index.id IN ({', '.join(id_params)}))" - ) - - if not branches: - return "1 = 0" - return f"({' OR '.join(branches)})" - - async def purge_stale_search_index_rows( session_maker: async_sessionmaker[AsyncSession], project_id: int, @@ -390,6 +240,8 @@ class SearchRepositoryBase(ABC): _vector_tables_initialized: bool _semantic_vector_index: SemanticVectorIndex _semantic_vector_index_name: str = "" + # Runs compiled full-text statements for this repository's engine. + _fts: FtsBackend def __init__(self, session_maker: async_sessionmaker[AsyncSession], project_id: int): """Initialize with session maker and project_id filter. @@ -458,7 +310,6 @@ async def init_search_index(self) -> None: """ pass - @abstractmethod async def search( self, search_text: Optional[str] = None, @@ -477,42 +328,71 @@ async def search( limit: int = 10, offset: int = 0, allow_relaxed: bool = False, + session: AsyncSession | None = None, *, candidate_keys: Sequence[SearchIndexKey] | None = None, trace: SearchTraceCollector | None = None, ) -> List[SearchIndexRow]: - """Search across all indexed content. + """Search this repository's project. - Args: - search_text: Full-text search across title and content - permalink: Exact permalink match - permalink_match: Permalink pattern match (supports *) - title: Title search - note_types: Filter by note types (from metadata.note_type) - after_date: Filter by created_at > after_date - search_item_types: Filter by SearchItemType (ENTITY, OBSERVATION, RELATION) - categories: Filter observations by exact category (e.g. "requirement") - metadata_filters: Structured frontmatter metadata filters - file_path_prefix: Directory subtree scope, matched against file_path - temporal: Authored valid-time filter. Unlike after_date, which reads the - note's edit bookkeeping, this reads the time an observation claims to - be true of the world. Sources without such a claim are excluded. - limit: Maximum results to return - offset: Number of results to skip - candidate_keys: Restrict results to these ``(type, id)`` search rows. ``None`` - searches the whole project; an empty sequence matches nothing. Honored by - the full-text pass, which is where vector and hybrid retrieval evaluate - their structured filters: that pass asks which of a known candidate set a - filter admits instead of paging the filter's whole match set (#1431). - - Returns: - List of SearchIndexRow results with relevance scores + ``candidate_keys`` restricts results to those ``(type, id)`` search rows. + ``None`` searches the whole scope; an empty sequence matches nothing. Honored by + the full-text pass, which is where vector and hybrid retrieval evaluate their + structured filters: that pass asks which of a known candidate set a filter + admits instead of paging the filter's whole match set (#1431). - Backend-specific implementations: - - SQLite: Uses MATCH operator and bm25() for scoring - - Postgres: Uses @@ operator and ts_rank() for scoring + ``allow_relaxed=True`` retries a zero-result strict multi-word query with + OR-joined content terms. Only the hybrid path opts in: its FTS branch otherwise + contributes nothing for question-form queries. """ - pass + # --- Vector and hybrid: shared retrieval over this repository's scope --- + dispatched = await self._dispatch_retrieval_mode( + search_text=search_text, + permalink=permalink, + permalink_match=permalink_match, + title=title, + note_types=note_types, + after_date=after_date, + search_item_types=search_item_types, + categories=categories, + metadata_filters=metadata_filters, + file_path_prefix=file_path_prefix, + temporal=temporal, + retrieval_mode=retrieval_mode, + min_similarity=min_similarity, + limit=limit, + offset=offset, + trace=trace, + ) + if dispatched is not None: + return dispatched + + # --- Full text: the engine runs the compiled statement --- + query = PreparedSearchQuery( + search_text=search_text, + permalink=permalink, + permalink_match=permalink_match, + title=title, + note_types=note_types, + search_item_types=search_item_types, + categories=categories, + after_date=after_date, + metadata_filters=metadata_filters, + file_path_prefix=file_path_prefix, + temporal=temporal, + retrieval_mode=retrieval_mode, + min_similarity=min_similarity, + ) + return await self._fts.search( + self.scope, + query, + limit=limit, + offset=offset, + allow_relaxed=allow_relaxed, + session=session, + candidate_keys=candidate_keys, + trace=trace, + ) async def count( self, @@ -531,10 +411,25 @@ async def count( min_similarity: Optional[float] = None, allow_relaxed: bool = False, ) -> int: - """Count results when a backend-specific COUNT query is available.""" + """Count full-text matches with the same filters as ``search``.""" if retrieval_mode != SearchRetrievalMode.FTS: raise ValueError("Exact counts are only supported for full-text search retrieval.") - raise NotImplementedError("Backend search repositories must implement full-text counts.") + query = PreparedSearchQuery( + search_text=search_text, + permalink=permalink, + permalink_match=permalink_match, + title=title, + note_types=note_types, + search_item_types=search_item_types, + categories=categories, + after_date=after_date, + metadata_filters=metadata_filters, + file_path_prefix=file_path_prefix, + temporal=temporal, + retrieval_mode=retrieval_mode, + min_similarity=min_similarity, + ) + return await self._fts.count(self.scope, query, allow_relaxed=allow_relaxed) # ------------------------------------------------------------------ # Abstract methods — semantic search (backend-specific DB operations) @@ -573,7 +468,7 @@ async def _run_vector_query( if trace is not None: trace.readiness = await read_manifest_readiness( session, - self.project_id, + self.scope, self._semantic_vector_index_name, self._embedding_model_key(), ) @@ -588,7 +483,7 @@ async def _run_vector_query( if trace is not None and trace.readiness is None: trace.readiness = await read_manifest_readiness( session, - self.project_id, + self.scope, self._semantic_vector_index_name, self._embedding_model_key(), ) @@ -715,11 +610,11 @@ async def _hydrate_vector_matches( chunks_by_key: dict[VectorKey, str] = {} for batch_start in range(0, len(matches), VECTOR_HYDRATION_BATCH_SIZE): batch = matches[batch_start : batch_start + VECTOR_HYDRATION_BATCH_SIZE] - params: dict[str, object] = { - "project_id": self.project_id, + params: dict[str, Any] = { "vector_index": self._semantic_vector_index_name, "embedding_model": self._embedding_model_key(), } + manifest_predicate = current_vector_manifest_predicate(self.scope, params) predicates: list[str] = [] for index, match in enumerate(batch): params[f"entity_id_{index}"] = match.key.entity_id @@ -734,7 +629,7 @@ async def _hydrate_vector_matches( result = await session.execute( text( "SELECT entity_id, chunk_key, chunk_text FROM search_vector_chunks " - "WHERE " + CURRENT_VECTOR_MANIFEST_PREDICATE + " " + "WHERE " + manifest_predicate + " " "AND (" + " OR ".join(predicates) + ")" ), params, @@ -770,7 +665,7 @@ async def _hydrate_vector_matches( for match in matches if match.key not in chunks_by_key ] - drops = await classify_hydration_drops(session, self.project_id, dropped_keys) + drops = await classify_hydration_drops(session, self.scope, dropped_keys) chunk_matches: dict[SearchIndexKey, list[tuple[str, float, int | None]]] = {} malformed_drops: list[HydrationDropped] = [] for row in hydrated: @@ -2796,17 +2691,15 @@ async def _fetch_search_index_rows_by_ids( if not row_ids: return {} placeholders = ",".join(f":id_{idx}" for idx in range(len(row_ids))) - params: dict[str, Any] = { - **{f"id_{idx}": rid for idx, rid in enumerate(row_ids)}, - "project_id": self.project_id, - } + params: dict[str, Any] = {f"id_{idx}": rid for idx, rid in enumerate(row_ids)} + scope_predicate = self.scope.predicate("project_id", params) sql = f""" SELECT project_id, id, title, permalink, file_path, type, metadata, from_id, to_id, relation_type, entity_id, content_snippet, category, created_at, updated_at, 0 as score FROM search_index - WHERE project_id = :project_id + WHERE {scope_predicate} AND id IN ({placeholders}) """ result: dict[SearchIndexKey, SearchIndexRow] = {} diff --git a/src/basic_memory/repository/search_trace.py b/src/basic_memory/repository/search_trace.py index eb7f53ffa..c090b08b0 100644 --- a/src/basic_memory/repository/search_trace.py +++ b/src/basic_memory/repository/search_trace.py @@ -1,5 +1,6 @@ """Typed, execution-native trace values for the search retrieval pipeline.""" +from basic_memory.repository.search_scope import ProjectScope from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Any, Literal @@ -508,25 +509,23 @@ class HydrationDropKey: async def read_manifest_readiness( session: Any, - project_id: int, + scope: ProjectScope, vector_index: str, embedding_model: str, ) -> ManifestReadiness: """Count configured readiness and rows stored under another vector identity.""" from sqlalchemy import text + params: dict[str, Any] = {"vector_index": vector_index, "embedding_model": embedding_model} + scope_predicate = scope.predicate("project_id", params) readiness_result = await session.execute( text( "SELECT embedding_status, COUNT(*) AS row_count " - "FROM search_vector_chunks WHERE project_id = :project_id " + f"FROM search_vector_chunks WHERE {scope_predicate} " "AND vector_index = :vector_index AND embedding_model = :embedding_model " "GROUP BY embedding_status" ), - { - "project_id": project_id, - "vector_index": vector_index, - "embedding_model": embedding_model, - }, + params, ) counts = { str(row["embedding_status"]): int(row["row_count"]) @@ -534,14 +533,10 @@ async def read_manifest_readiness( } other_result = await session.execute( text( - "SELECT COUNT(*) FROM search_vector_chunks WHERE project_id = :project_id " + f"SELECT COUNT(*) FROM search_vector_chunks WHERE {scope_predicate} " "AND (vector_index <> :vector_index OR embedding_model <> :embedding_model)" ), - { - "project_id": project_id, - "vector_index": vector_index, - "embedding_model": embedding_model, - }, + params, ) return ManifestReadiness( configured_index=vector_index, @@ -554,7 +549,7 @@ async def read_manifest_readiness( async def classify_hydration_drops( session: Any, - project_id: int, + scope: ProjectScope, dropped_keys: Sequence[HydrationDropKey], ) -> tuple[HydrationDropped, ...]: """Classify adapter hits rejected by authoritative manifest hydration.""" @@ -570,7 +565,8 @@ async def classify_hydration_drops( HYDRATION_DROP_CLASSIFICATION_BATCH_SIZE, ): batch = dropped_keys[batch_start : batch_start + HYDRATION_DROP_CLASSIFICATION_BATCH_SIZE] - params: dict[str, object] = {"project_id": project_id} + params: dict[str, Any] = {} + scope_predicate = scope.predicate("project_id", params) predicates: list[str] = [] for index, dropped in enumerate(batch): params[f"entity_id_{index}"] = dropped.entity_id @@ -585,7 +581,7 @@ async def classify_hydration_drops( result = await session.execute( text( "SELECT entity_id, chunk_key, embedding_model, vector_index, embedding_status " - "FROM search_vector_chunks WHERE project_id = :project_id AND (" + f"FROM search_vector_chunks WHERE {scope_predicate} AND (" + " OR ".join(predicates) + ")" ), diff --git a/src/basic_memory/repository/sqlite_search_query.py b/src/basic_memory/repository/sqlite_search_query.py index 9fa161bab..42dfdbd67 100644 --- a/src/basic_memory/repository/sqlite_search_query.py +++ b/src/basic_memory/repository/sqlite_search_query.py @@ -1,32 +1,35 @@ -"""SQLite FTS5 query preparation: term syntax and filter compilation. +"""SQLite FTS5 query preparation and execution. -Pure functions over a ``ProjectScope`` and the caller's filters. Nothing here opens -a session or owns an index; the repository resolves the one piece of live schema -state the compiler needs (the entity table's columns) and passes it in. +Term preparation and filter compilation are pure functions over a ``ProjectScope`` and +a ``PreparedSearchQuery``. ``SQLiteFts`` runs the compiled statement and owns FTS5's +failure semantics. Nothing here initializes or mutates an index. """ import re +import time from collections.abc import Collection, Sequence -from datetime import datetime from typing import Any +import logfire +from loguru import logger +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from basic_memory import db from basic_memory.repository.metadata_filters import build_sqlite_json_path, parse_metadata_filters from basic_memory.repository.script_ngrams import analyze_script_query from basic_memory.repository.search_filters import ( AFTER_DATE_ORDER_BY, SQLITE_FILTER_DIALECT, CompiledFilter, - shared_filter_conditions, -) -from basic_memory.repository.search_query import relaxed_query_words -from basic_memory.repository.search_repository_base import ( - SearchIndexKey, metadata_contains_like_condition, metadata_filter_content_type_condition, + shared_filter_conditions, ) +from basic_memory.repository.search_index_row import SearchIndexKey, SearchIndexRow +from basic_memory.repository.search_query import PreparedSearchQuery, relaxed_query_words from basic_memory.repository.search_scope import ProjectScope -from basic_memory.schemas.search import SearchItemType -from basic_memory.temporal import TemporalFilter +from basic_memory.repository.search_trace import SearchTraceCollector, build_fts_page_stage SQLITE_WORD_COLUMNS = "{title content_stems content_snippet}" @@ -39,6 +42,24 @@ _SPACE_OR_SPECIAL_CHARS = frozenset(" .:;,<>?/-") _BOOLEAN_OPERATOR_PATTERN = r"(\bAND\b|\bOR\b|\bNOT\b)" +# Every FTS statement returns these columns plus a score. +_RESULT_COLUMNS = """ + search_index.project_id, + search_index.id, + search_index.title, + search_index.permalink, + search_index.file_path, + search_index.type, + search_index.metadata, + search_index.from_id, + search_index.to_id, + search_index.relation_type, + search_index.entity_id, + search_index.content_snippet, + search_index.category, + search_index.created_at, + search_index.updated_at""" + # --- Term preparation --- @@ -188,19 +209,9 @@ def is_fts5_syntax_error(exc: Exception) -> bool: def compile_fts_filter( scope: ProjectScope, + query: PreparedSearchQuery, *, entity_columns: Collection[str], - search_text: str | None = None, - permalink: str | None = None, - permalink_match: str | None = None, - title: str | None = None, - note_types: Sequence[str] | None = None, - after_date: datetime | None = None, - search_item_types: Sequence[SearchItemType] | None = None, - categories: Sequence[str] | None = None, - metadata_filters: dict[str, Any] | None = None, - file_path_prefix: str | None = None, - temporal: TemporalFilter | None = None, candidate_keys: Sequence[SearchIndexKey] | None = None, ) -> CompiledFilter: """Compile SQLite FTS FROM/WHERE/score shared by search and count. @@ -210,22 +221,13 @@ def compile_fts_filter( """ params: dict[str, Any] = {} conditions = shared_filter_conditions( - scope, - params, - dialect=SQLITE_FILTER_DIALECT, - permalink=permalink, - file_path_prefix=file_path_prefix, - candidate_keys=candidate_keys, - search_item_types=search_item_types, - categories=categories, - note_types=note_types, - after_date=after_date, - temporal=temporal, + scope, params, dialect=SQLITE_FILTER_DIALECT, query=query, candidate_keys=candidate_keys ) match_conditions: list[str] = [] from_clause = "search_index" score_expression = "bm25(search_index)" preserve_match_score = False + search_text = query.search_text # Wildcard-only and blank text add no text condition: every row matches. if search_text and search_text.strip() not in ("", "*"): @@ -268,15 +270,15 @@ def compile_fts_filter( "search_index.content_snippet MATCH :text)" ) - if title: - params["title_text"] = prepare_search_term(title.strip(), is_prefix=False) + if query.title: + params["title_text"] = prepare_search_term(query.title.strip(), is_prefix=False) match_conditions.append("search_index.title MATCH :title_text") - if permalink_match: + if query.permalink_match: # GLOB patterns keep their syntax; prepare_search_term would quote the slashes. - permalink_text = permalink_match.lower().strip() + permalink_text = query.permalink_match.lower().strip() params["permalink"] = permalink_text - if "*" in permalink_match: + if "*" in query.permalink_match: conditions.append("search_index.permalink GLOB :permalink") elif "/" in permalink_text: conditions.append("search_index.permalink = :permalink") @@ -285,8 +287,8 @@ def compile_fts_filter( params["permalink"] = prepare_search_term(permalink_text, is_prefix=False) match_conditions.append("search_index.permalink MATCH :permalink") - if metadata_filters: - parsed_filters = parse_metadata_filters(metadata_filters) + if query.metadata_filters: + parsed_filters = parse_metadata_filters(query.metadata_filters) from_clause = "search_index JOIN entity ON search_index.entity_id = entity.id" # Frontmatter filters answer for notes only; see # metadata_filter_content_type_condition for why every regular file would @@ -397,7 +399,7 @@ def compile_fts_filter( # word-column OR predicates cannot evaluate bm25 in the same derived query. # Outcome: rank script matches before joining metadata; retain the established # rowid-filter path for word-only searches. - if metadata_filters and match_conditions: + if query.metadata_filters and match_conditions: match_where = " AND ".join(match_conditions) if preserve_match_score: from_clause = ( @@ -418,6 +420,198 @@ def compile_fts_filter( from_clause=from_clause, where_clause=" AND ".join(conditions), params=params, - order_by_clause=AFTER_DATE_ORDER_BY if after_date else "", + order_by_clause=AFTER_DATE_ORDER_BY if query.after_date else "", score_expression=score_expression, ) + + +# --- Execution --- + + +class SQLiteFts: + """Run FTS5 statements for any scope in one database. + + Holds the one piece of live schema state compilation needs: the entity table's + columns, read once per instance so generated frontmatter columns are used when the + database has them. + """ + + def __init__(self, session_maker: async_sessionmaker[AsyncSession]) -> None: + self._session_maker = session_maker + self._entity_columns: frozenset[str] | None = None + + async def _entity_column_names(self) -> frozenset[str]: + if self._entity_columns is None: + async with db.scoped_session(self._session_maker) as session: + result = await session.execute(text("PRAGMA table_info(entity)")) + self._entity_columns = frozenset(row[1] for row in result.fetchall()) + return self._entity_columns + + async def search( + self, + scope: ProjectScope, + query: PreparedSearchQuery, + *, + limit: int, + offset: int, + allow_relaxed: bool = False, + session: AsyncSession | None = None, + candidate_keys: Sequence[SearchIndexKey] | None = None, + trace: SearchTraceCollector | None = None, + ) -> list[SearchIndexRow]: + """Run one FTS5 page. + + ``allow_relaxed=True`` retries a zero-result strict multi-word query with + OR-joined content terms. Only the hybrid path opts in: its FTS branch otherwise + contributes nothing for question-form queries. Service-level FTS searches keep + their own conservative fallback. + """ + search_text = query.search_text + # Generated frontmatter columns are read only when a metadata filter needs them. + entity_columns = ( + await self._entity_column_names() if query.metadata_filters else frozenset() + ) + compiled = compile_fts_filter( + scope, query, entity_columns=entity_columns, candidate_keys=candidate_keys + ) + params = compiled.params + params["limit"] = limit + params["offset"] = offset + relaxed_search_text = search_text + if search_text and "script_text" in params: + relaxed_search_text = analyze_script_query(search_text.strip()).word_text + + sql = f""" + SELECT{_RESULT_COLUMNS}, + {compiled.score_expression} as score + FROM {compiled.from_clause} + WHERE {compiled.where_clause} + ORDER BY score ASC {compiled.order_by_clause} + LIMIT :limit + OFFSET :offset + """ + + logger.trace(f"Search {sql} params: {params}") + fts_started_at = time.perf_counter() if trace is not None else None + + async def run_search(active_session: AsyncSession): + result = await active_session.execute(text(sql), params) + rows = result.fetchall() + relaxed_fallback_used = False + # Trigger: multi-word natural-language query matched nothing under the + # default all-terms-AND semantics. + # Why: questions ("when did X do Y") rarely have every word in one + # document; without relaxation the FTS half of hybrid search contributes + # zero candidates and ranking degrades to vector-only. + # Outcome: one retry with OR-joined prefix terms; bm25 still ranks + # multi-term matches first. + relaxed = relaxed_fts_text(relaxed_search_text) if allow_relaxed and not rows else None + if relaxed and params.get("text"): + relaxed_fallback_used = True + params["text"] = ( + f"{SQLITE_WORD_COLUMNS}: ({relaxed})" if "script_text" in params else relaxed + ) + logger.debug( + "Strict SQLite FTS returned 0 results; retrying relaxed FTS query " + f"strict='{search_text}' relaxed='{relaxed}'" + ) + with logfire.span( + "search.relaxed_fts_retry", + backend="sqlite", + token_count=len(relaxed_query_words(relaxed_search_text) or ()), + limit=limit, + offset=offset, + ): + result = await active_session.execute(text(sql), params) + rows = result.fetchall() + return rows, relaxed_fallback_used + + try: + if session is not None: + rows, relaxed_fallback_used = await run_search(session) + else: + async with db.scoped_session(self._session_maker) as owned_session: + rows, relaxed_fallback_used = await run_search(owned_session) + except Exception as e: + # An FTS5 syntax error answers with no rows rather than failing the request. + if is_fts5_syntax_error(e): # pragma: no cover + logger.warning(f"FTS5 syntax error for search term: {search_text}, error: {e}") + if trace is not None: + trace.fts = build_fts_page_stage( + [], + relaxed_fallback_used=False, + fts_ms=( + (time.perf_counter() - fts_started_at) * 1000 + if fts_started_at is not None + else None + ), + ) + return [] + logger.error(f"Database error during search: {e}") + raise + + results = [SearchIndexRow.from_mapping(row._asdict()) for row in rows] + if trace is not None: + trace.fts = build_fts_page_stage( + [((row.type, row.id), row.score or 0.0) for row in results], + relaxed_fallback_used=relaxed_fallback_used, + fts_ms=( + (time.perf_counter() - fts_started_at) * 1000 + if fts_started_at is not None + else None + ), + ) + + logger.trace(f"Found {len(results)} search results") + for r in results: + logger.trace( + f"Search result: project_id: {r.project_id} type:{r.type} title: {r.title} permalink: {r.permalink} score: {r.score}" + ) + return results + + async def count( + self, + scope: ProjectScope, + query: PreparedSearchQuery, + *, + allow_relaxed: bool = False, + ) -> int: + """Count rows matching the FTS5 query, with the same relaxed retry as search.""" + search_text = query.search_text + entity_columns = ( + await self._entity_column_names() if query.metadata_filters else frozenset() + ) + compiled = compile_fts_filter(scope, query, entity_columns=entity_columns) + params = compiled.params + sql = f"SELECT COUNT(*) FROM {compiled.from_clause} WHERE {compiled.where_clause}" + logger.trace(f"Count {sql} params: {params}") + relaxed_search_text = search_text + if search_text and "script_text" in params: + relaxed_search_text = analyze_script_query(search_text.strip()).word_text + try: + async with db.scoped_session(self._session_maker) as session: + result = await session.execute(text(sql), params) + total = int(result.scalar_one()) + relaxed = ( + relaxed_fts_text(relaxed_search_text) if allow_relaxed and total == 0 else None + ) + if relaxed and params.get("text"): + params["text"] = ( + f"{SQLITE_WORD_COLUMNS}: ({relaxed})" + if "script_text" in params + else relaxed + ) + with logfire.span( + "search.count.relaxed_fts_retry", + backend="sqlite", + token_count=len(relaxed_query_words(relaxed_search_text) or ()), + ): + result = await session.execute(text(sql), params) + total = int(result.scalar_one()) + return total + except Exception as e: + if is_fts5_syntax_error(e): # pragma: no cover + logger.warning(f"FTS5 syntax error for search term: {search_text}, error: {e}") + return 0 + logger.error(f"Database error during search count: {e}") + raise diff --git a/src/basic_memory/repository/sqlite_search_repository.py b/src/basic_memory/repository/sqlite_search_repository.py index 11bb39ff0..aa75461e7 100644 --- a/src/basic_memory/repository/sqlite_search_repository.py +++ b/src/basic_memory/repository/sqlite_search_repository.py @@ -1,13 +1,10 @@ """SQLite FTS5-based search repository implementation.""" import asyncio -import time from collections.abc import Sequence from contextlib import asynccontextmanager -from datetime import datetime -from typing import Any, override, List, Optional +from typing import override, List -import logfire from loguru import logger from sqlalchemy import text from sqlalchemy.exc import OperationalError as SAOperationalError @@ -26,29 +23,15 @@ from basic_memory.repository.rerank_provider import RerankProvider from basic_memory.repository.rerank_provider_factory import create_rerank_provider from basic_memory.repository.search_index_row import SearchIndexRow -from basic_memory.repository.search_query import relaxed_query_words from basic_memory.repository.search_repository_base import ( - SearchIndexKey, SearchRepositoryBase, ) -from basic_memory.repository.script_ngrams import analyze_script_query -from basic_memory.repository.search_trace import ( - SearchTraceCollector, - build_fts_page_stage, -) -from basic_memory.repository.sqlite_search_query import ( - SQLITE_WORD_COLUMNS, - compile_fts_filter, - is_fts5_syntax_error, - relaxed_fts_text, -) +from basic_memory.repository.sqlite_search_query import SQLiteFts from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError from basic_memory.repository.semantic_vector_index import SemanticVectorIndex from basic_memory.repository.semantic_vector_sync import StagedVectorDeletion from basic_memory.repository.semantic_vector_index_factory import build_vector_index_scope from basic_memory.repository.sqlite_vec_index import SQLiteVecIndex -from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode -from basic_memory.temporal import TemporalFilter class SQLiteSearchRepository(SearchRepositoryBase): @@ -72,7 +55,7 @@ def __init__( rerank_provider: RerankProvider | None = None, ): super().__init__(session_maker, project_id) - self._entity_columns: set[str] | None = None + self._fts = SQLiteFts(session_maker) self._app_config = app_config or ConfigManager().config self._semantic_enabled = self._app_config.semantic_search_enabled self._semantic_vector_k = self._app_config.semantic_vector_k @@ -109,13 +92,6 @@ def __init__( ), ) - async def _get_entity_columns(self) -> set[str]: - if self._entity_columns is None: - async with db.scoped_session(self.session_maker) as session: - result = await session.execute(text("PRAGMA table_info(entity)")) - self._entity_columns = {row[1] for row in result.fetchall()} - return self._entity_columns - @override async def init_search_index(self): """Create FTS5 virtual table for search if it doesn't exist. @@ -358,25 +334,6 @@ async def _prepare_vector_session(self, session: AsyncSession) -> None: """Load sqlite-vec extension for the session.""" await self._ensure_sqlite_vec_loaded(session) - # sqlite-vec hard limit for knn k parameter - SQLITE_VEC_MAX_K = 4096 - - @override - async def _run_vector_query( - self, - session: AsyncSession, - query_embedding: list[float], - candidate_limit: int, - *, - trace: SearchTraceCollector | None = None, - ) -> list[dict[str, Any]]: - return await super()._run_vector_query( - session, - query_embedding, - candidate_limit, - trace=trace, - ) - @override async def _delete_entity_chunks( self, @@ -492,280 +449,3 @@ async def index_item(self, search_index_row: SearchIndexRow) -> None: async def bulk_index_items(self, search_index_rows: List[SearchIndexRow]) -> None: """Index multiple rows in FTS only.""" await super().bulk_index_items(search_index_rows) - - # ------------------------------------------------------------------ - # FTS search (backend-specific) - # ------------------------------------------------------------------ - - @override - async def search( - self, - search_text: Optional[str] = None, - permalink: Optional[str] = None, - permalink_match: Optional[str] = None, - title: Optional[str] = None, - note_types: Optional[List[str]] = None, - after_date: Optional[datetime] = None, - search_item_types: Optional[List[SearchItemType]] = None, - categories: Optional[List[str]] = None, - metadata_filters: Optional[dict[str, Any]] = None, - file_path_prefix: Optional[str] = None, - temporal: Optional[TemporalFilter] = None, - retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, - min_similarity: Optional[float] = None, - limit: int = 10, - offset: int = 0, - allow_relaxed: bool = False, - session: AsyncSession | None = None, - *, - candidate_keys: Sequence[SearchIndexKey] | None = None, - trace: SearchTraceCollector | None = None, - ) -> List[SearchIndexRow]: - """Search across all indexed content using SQLite FTS5. - - ``allow_relaxed=True`` retries a zero-result strict multi-word query - with OR-joined content terms. Only the hybrid path opts in: its FTS - branch otherwise contributes nothing for question-form queries. - Service-level FTS searches keep their own conservative fallback. - """ - # --- Dispatch vector / hybrid modes (shared logic) --- - dispatched = await self._dispatch_retrieval_mode( - search_text=search_text, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - retrieval_mode=retrieval_mode, - min_similarity=min_similarity, - limit=limit, - offset=offset, - trace=trace, - ) - if dispatched is not None: - return dispatched - - # --- FTS mode (SQLite-specific) --- - # Generated frontmatter columns are read only when a metadata filter needs them. - entity_columns = await self._get_entity_columns() if metadata_filters else frozenset() - compiled = compile_fts_filter( - self.scope, - entity_columns=entity_columns, - search_text=search_text, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - candidate_keys=candidate_keys, - ) - params = compiled.params - params["limit"] = limit - params["offset"] = offset - relaxed_search_text = search_text - if search_text and "script_text" in params: - relaxed_search_text = analyze_script_query(search_text.strip()).word_text - - sql = f""" - SELECT - search_index.project_id, - search_index.id, - search_index.title, - search_index.permalink, - search_index.file_path, - search_index.type, - search_index.metadata, - search_index.from_id, - search_index.to_id, - search_index.relation_type, - search_index.entity_id, - search_index.content_snippet, - search_index.category, - search_index.created_at, - search_index.updated_at, - {compiled.score_expression} as score - FROM {compiled.from_clause} - WHERE {compiled.where_clause} - ORDER BY score ASC {compiled.order_by_clause} - LIMIT :limit - OFFSET :offset - """ - - logger.trace(f"Search {sql} params: {params}") - fts_started_at = time.perf_counter() if trace is not None else None - - async def run_search(active_session: AsyncSession): - result = await active_session.execute(text(sql), params) - rows = result.fetchall() - relaxed_fallback_used = False - # Trigger: multi-word natural-language query matched nothing - # under the default all-terms-AND semantics. - # Why: questions ("when did X do Y") rarely have every word in - # one document; without relaxation the FTS half of hybrid - # search contributes zero candidates and ranking degrades to - # vector-only. - # Outcome: one retry with OR-joined prefix terms; bm25 still - # ranks multi-term matches first. - relaxed = relaxed_fts_text(relaxed_search_text) if allow_relaxed and not rows else None - if relaxed and params.get("text"): - relaxed_fallback_used = True - params["text"] = ( - f"{SQLITE_WORD_COLUMNS}: ({relaxed})" if "script_text" in params else relaxed - ) - logger.debug( - "Strict SQLite FTS returned 0 results; retrying relaxed FTS query " - f"strict='{search_text}' relaxed='{relaxed}'" - ) - with logfire.span( - "search.relaxed_fts_retry", - backend="sqlite", - token_count=len(relaxed_query_words(relaxed_search_text) or ()), - limit=limit, - offset=offset, - ): - result = await active_session.execute(text(sql), params) - rows = result.fetchall() - return rows, relaxed_fallback_used - - try: - if session is not None: - rows, relaxed_fallback_used = await run_search(session) - else: - async with db.scoped_session(self.session_maker) as owned_session: - rows, relaxed_fallback_used = await run_search(owned_session) - except Exception as e: - # Handle FTS5 syntax errors and provide user-friendly feedback - if is_fts5_syntax_error(e): # pragma: no cover - logger.warning(f"FTS5 syntax error for search term: {search_text}, error: {e}") - # Return empty results rather than crashing - if trace is not None: - trace.fts = build_fts_page_stage( - [], - relaxed_fallback_used=False, - fts_ms=( - (time.perf_counter() - fts_started_at) * 1000 - if fts_started_at is not None - else None - ), - ) - return [] - else: - # Re-raise other database errors - logger.error(f"Database error during search: {e}") - raise - - results = [SearchIndexRow.from_mapping(row._asdict()) for row in rows] - if trace is not None: - trace.fts = build_fts_page_stage( - [((row.type, row.id), row.score or 0.0) for row in results], - relaxed_fallback_used=relaxed_fallback_used, - fts_ms=( - (time.perf_counter() - fts_started_at) * 1000 - if fts_started_at is not None - else None - ), - ) - - logger.trace(f"Found {len(results)} search results") - for r in results: - logger.trace( - f"Search result: project_id: {r.project_id} type:{r.type} title: {r.title} permalink: {r.permalink} score: {r.score}" - ) - - return results - - @override - async def count( - self, - search_text: Optional[str] = None, - permalink: Optional[str] = None, - permalink_match: Optional[str] = None, - title: Optional[str] = None, - note_types: Optional[List[str]] = None, - after_date: Optional[datetime] = None, - search_item_types: Optional[List[SearchItemType]] = None, - categories: Optional[List[str]] = None, - metadata_filters: Optional[dict[str, Any]] = None, - file_path_prefix: Optional[str] = None, - temporal: Optional[TemporalFilter] = None, - retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, - min_similarity: Optional[float] = None, - allow_relaxed: bool = False, - ) -> int: - """Count indexed content matching the SQLite FTS query.""" - if retrieval_mode != SearchRetrievalMode.FTS: - return await super().count( - search_text=search_text, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - retrieval_mode=retrieval_mode, - min_similarity=min_similarity, - ) - - entity_columns = await self._get_entity_columns() if metadata_filters else frozenset() - compiled = compile_fts_filter( - self.scope, - entity_columns=entity_columns, - search_text=search_text, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - ) - params = compiled.params - sql = f"SELECT COUNT(*) FROM {compiled.from_clause} WHERE {compiled.where_clause}" - logger.trace(f"Count {sql} params: {params}") - relaxed_search_text = search_text - if search_text and "script_text" in params: - relaxed_search_text = analyze_script_query(search_text.strip()).word_text - try: - async with db.scoped_session(self.session_maker) as session: - result = await session.execute(text(sql), params) - total = int(result.scalar_one()) - relaxed = ( - relaxed_fts_text(relaxed_search_text) if allow_relaxed and total == 0 else None - ) - if relaxed and params.get("text"): - params["text"] = ( - f"{SQLITE_WORD_COLUMNS}: ({relaxed})" - if "script_text" in params - else relaxed - ) - with logfire.span( - "search.count.relaxed_fts_retry", - backend="sqlite", - token_count=len(relaxed_query_words(relaxed_search_text) or ()), - ): - result = await session.execute(text(sql), params) - total = int(result.scalar_one()) - return total - except Exception as e: - if is_fts5_syntax_error(e): # pragma: no cover - logger.warning(f"FTS5 syntax error for search term: {search_text}, error: {e}") - return 0 - logger.error(f"Database error during search count: {e}") - raise diff --git a/src/basic_memory/services/project_readiness.py b/src/basic_memory/services/project_readiness.py index f37a14fc2..ed7bc1ca8 100644 --- a/src/basic_memory/services/project_readiness.py +++ b/src/basic_memory/services/project_readiness.py @@ -24,7 +24,8 @@ from basic_memory.repository.embedding_provider_factory import ( configured_embedding_provider_identity, ) -from basic_memory.repository.search_repository_base import CURRENT_VECTOR_MANIFEST_PREDICATE +from basic_memory.repository.search_repository_base import current_vector_manifest_predicate +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_vector_index_factory import ( resolve_semantic_vector_index_name, ) @@ -310,10 +311,21 @@ async def _embedding_counts( # embedded -- the third time a count and the thing it measures disagreed # (#1440 review). The marker is written by the sharded sync itself, in # `record_entity_vector_deferrals`, so the two cannot drift. + manifest_params: dict[str, object] = { + "project_id": project_id, + "vector_index": resolve_semantic_vector_index_name( + self.app_config, + self.app_config.database_backend, + ), + "embedding_model": configured_embedding_provider_identity(self.app_config), + } + manifest_predicate = current_vector_manifest_predicate( + ProjectScope.single(project_id), manifest_params + ) usable_result = await session.execute( text( "SELECT DISTINCT entity_id FROM search_vector_chunks " - "WHERE " + CURRENT_VECTOR_MANIFEST_PREDICATE + " " + "WHERE " + manifest_predicate + " " # Applied as a subquery so the shared predicate is used verbatim # rather than rewritten to carry a table alias. "AND entity_id NOT IN (" @@ -321,14 +333,7 @@ async def _embedding_counts( " AND vector_sync_deferred_at IS NOT NULL" ")" ), - { - "project_id": project_id, - "vector_index": resolve_semantic_vector_index_name( - self.app_config, - self.app_config.database_backend, - ), - "embedding_model": configured_embedding_provider_identity(self.app_config), - }, + manifest_params, ) usable_entity_ids = {int(entity_id) for entity_id in usable_result.scalars().all()} return len(owed_entity_ids), len(owed_entity_ids & usable_entity_ids) diff --git a/src/basic_memory/services/search_service.py b/src/basic_memory/services/search_service.py index 301e643f2..0a0f70317 100644 --- a/src/basic_memory/services/search_service.py +++ b/src/basic_memory/services/search_service.py @@ -4,7 +4,7 @@ import ast import re from collections.abc import Mapping, Sequence -from dataclasses import dataclass, replace +from dataclasses import replace from datetime import datetime from typing import Any, List, Optional, Set, Dict @@ -23,7 +23,7 @@ SearchIndexRow, SearchRepository, ) -from basic_memory.repository.search_query import relaxed_query_words +from basic_memory.repository.search_query import PreparedSearchQuery, relaxed_query_words from basic_memory.repository.search_trace import SearchTraceCollector from basic_memory.schemas.base import normalize_note_type from basic_memory.schemas.search import SearchQuery, SearchItemType, SearchRetrievalMode @@ -42,25 +42,6 @@ MAX_CONTENT_STEMS_SIZE = 6000 -@dataclass(frozen=True) -class PreparedSearchQuery: - """Normalized query inputs shared by search and count.""" - - search_text: str | None - permalink: str | None - permalink_match: str | None - title: str | None - note_types: list[str] | None - search_item_types: list[SearchItemType] | None - categories: list[str] | None - after_date: datetime | None - metadata_filters: dict[str, Any] | None - file_path_prefix: str | None - temporal: TemporalFilter | None - retrieval_mode: SearchRetrievalMode - min_similarity: float | None - - def entity_embeddings_enabled(entity: Entity) -> bool: """Return whether semantic embeddings should be generated for this entity. diff --git a/tests/repository/test_hybrid_fusion.py b/tests/repository/test_hybrid_fusion.py index 4bf7135b1..0ad16896b 100644 --- a/tests/repository/test_hybrid_fusion.py +++ b/tests/repository/test_hybrid_fusion.py @@ -6,6 +6,8 @@ 3. Produces zero fused score when the source score is zero """ +from sqlalchemy.ext.asyncio import AsyncSession +from basic_memory.repository.search_scope import ProjectScope from collections.abc import Sequence from dataclasses import dataclass from datetime import datetime @@ -62,6 +64,7 @@ def __init__(self): self._vector_tables_initialized = True self.session_maker = None self.project_id = 1 + self.scope = ProjectScope.single(1) @override async def init_search_index(self): @@ -90,6 +93,7 @@ async def search( limit: int = 10, offset: int = 0, allow_relaxed: bool = False, + session: AsyncSession | None = None, *, candidate_keys: Sequence[SearchIndexKey] | None = None, trace: SearchTraceCollector | None = None, diff --git a/tests/repository/test_postgres_search_quoted_queries.py b/tests/repository/test_postgres_search_quoted_queries.py index 40a412619..3f59f562b 100644 --- a/tests/repository/test_postgres_search_quoted_queries.py +++ b/tests/repository/test_postgres_search_quoted_queries.py @@ -5,7 +5,7 @@ import pytest -import basic_memory.repository.postgres_search_repository as postgres_search_repository_module +import basic_memory.repository.postgres_search_query as postgres_search_query_module from basic_memory.repository.postgres_search_repository import PostgresSearchRepository from basic_memory.repository.search_index_row import SearchIndexRow @@ -61,7 +61,7 @@ async def test_quoted_or_phrases_complete_without_tsquery_recovery( ) syntax_errors: list[Exception] = [] - real_is_syntax_error = postgres_search_repository_module.is_tsquery_syntax_error + real_is_syntax_error = postgres_search_query_module.is_tsquery_syntax_error def record_syntax_error(exception: Exception) -> bool: is_syntax_error = real_is_syntax_error(exception) @@ -69,9 +69,9 @@ def record_syntax_error(exception: Exception) -> bool: syntax_errors.append(exception) return is_syntax_error - # The repository module binds the classifier at import; patch it where it is read. + # PostgresFts binds the classifier at import; patch it where it is read. monkeypatch.setattr( - postgres_search_repository_module, "is_tsquery_syntax_error", record_syntax_error + postgres_search_query_module, "is_tsquery_syntax_error", record_syntax_error ) query = '"incident response" OR "database recovery"' diff --git a/tests/repository/test_postgres_search_repository.py b/tests/repository/test_postgres_search_repository.py index 308509ac7..197571fd6 100644 --- a/tests/repository/test_postgres_search_repository.py +++ b/tests/repository/test_postgres_search_repository.py @@ -14,13 +14,13 @@ import basic_memory.repository.search_repository_base as search_repository_base_module from basic_memory.repository.litellm_provider import LiteLLMEmbeddingProvider import basic_memory.repository.postgres_search_query as postgres_search_query_module -import basic_memory.repository.postgres_search_repository as postgres_search_repository_module from basic_memory.repository.postgres_search_query import ( compile_fts_filter, prepare_search_term, prepare_single_term, relaxed_tsquery_text, ) +from basic_memory.repository.search_query import PreparedSearchQuery from basic_memory.repository.postgres_search_repository import ( PostgresSearchRepository, _strip_nul_from_row, @@ -276,19 +276,22 @@ async def test_postgres_search_repository_bulk_index_items_and_prepare_terms( assert prepare_single_term(" ") == " " assert prepare_single_term("coffee", is_prefix=False) == "coffee" - indexed = compile_fts_filter(repo.scope, search_text="coffee brewing", allow_relaxed=True) + indexed = compile_fts_filter( + repo.scope, PreparedSearchQuery(search_text="coffee brewing"), allow_relaxed=True + ) assert "FROM search_index AS candidate_parent" in indexed.from_clause assert "FROM search_index_fts_chunks AS candidate_chunk" in indexed.from_clause assert "querytree(to_tsquery('english', :text))" in indexed.from_clause assert indexed.params["text_candidate"] == "coffee:* | brewing:*" filtered = compile_fts_filter( - repo.scope, search_text="coffee brewing", metadata_filters={"status": "active"} + repo.scope, + PreparedSearchQuery(search_text="coffee brewing", metadata_filters={"status": "active"}), ) assert "AS fts_candidate" in filtered.from_clause assert "JOIN entity ON search_index.entity_id = entity.id" in filtered.from_clause - negated = compile_fts_filter(repo.scope, search_text="coffee NOT brewing") + negated = compile_fts_filter(repo.scope, PreparedSearchQuery(search_text="coffee NOT brewing")) assert "AS fts_candidate" in negated.from_clause assert "FROM search_index AS candidate_all" in negated.from_clause assert negated.params["text_candidate"] == "coffee | brewing" @@ -1273,7 +1276,7 @@ async def test_postgres_relaxes_after_strict_tsquery_syntax_error( ) syntax_errors: list[Exception] = [] - real_is_syntax_error = postgres_search_repository_module.is_tsquery_syntax_error + real_is_syntax_error = postgres_search_query_module.is_tsquery_syntax_error def record_syntax_error(exc: Exception) -> bool: is_syntax_error = real_is_syntax_error(exc) @@ -1281,9 +1284,9 @@ def record_syntax_error(exc: Exception) -> bool: syntax_errors.append(exc) return is_syntax_error - # The repository module binds the classifier at import; patch it where it is read. + # PostgresFts binds the classifier at import; patch it where it is read. monkeypatch.setattr( - postgres_search_repository_module, "is_tsquery_syntax_error", record_syntax_error + postgres_search_query_module, "is_tsquery_syntax_error", record_syntax_error ) query = "foo Date: Tue, 15 Sep 2026 18:55:34 -0500 Subject: [PATCH 3/8] refactor(core): extract SearchReader and SemanticSearch from the search repository (#1565) Signed-off-by: phernandez Co-authored-by: Claude Fable 5.1 --- CHANGELOG.md | 12 + .../repository/postgres_search_repository.py | 9 - src/basic_memory/repository/search_query.py | 22 + src/basic_memory/repository/search_reader.py | 1136 ++++++++++++++ .../repository/search_repository_base.py | 1330 +---------------- .../repository/sqlite_search_repository.py | 9 - .../services/project_readiness.py | 2 +- .../services/retrieval_inspect.py | 9 +- .../test_multilingual_benchmark_contract.py | 2 +- test-int/semantic/test_search_diagnostics.py | 15 +- test-int/semantic/test_semantic_coverage.py | 18 +- .../repository/test_distance_to_similarity.py | 24 - tests/repository/test_hybrid_fusion.py | 314 ++-- .../test_postgres_search_repository_unit.py | 2 +- tests/repository/test_rerank_pipeline.py | 253 ++-- .../test_search_file_path_prefix.py | 23 +- tests/repository/test_search_reader.py | 283 ++++ tests/repository/test_search_trace.py | 4 +- tests/repository/test_semantic_search_base.py | 63 +- tests/repository/test_semantic_vector_sync.py | 15 - ...est_vector_filter_candidate_restriction.py | 61 +- tests/repository/test_vector_pagination.py | 208 +-- .../repository/test_vector_temporal_filter.py | 132 +- tests/repository/test_vector_threshold.py | 512 ++----- tests/services/test_project_readiness.py | 2 +- 25 files changed, 2056 insertions(+), 2404 deletions(-) create mode 100644 src/basic_memory/repository/search_reader.py delete mode 100644 tests/repository/test_distance_to_similarity.py create mode 100644 tests/repository/test_search_reader.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a523a2777..8b57ec9a5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -131,6 +131,18 @@ layer with defaults, and the filter helpers both backends share move from the base into `search_filters`. No query behavior changes. +- **#1558**: Retrieval leaves `SearchRepositoryBase`. `SearchReader` runs one prepared + query over one `ProjectScope` in whichever mode it asks for, and `SemanticSearch` owns + vector and hybrid retrieval (adapter lookup, manifest hydration, the structured filter + pass, score fusion, reranking, pagination) over a `VectorRetrieval` that is present or + absent rather than probed with `hasattr`. The repository keeps what only it knows + (whether semantic search is enabled and its vector tables exist) and builds a reader + per call from its current state. Hydrated chunks are a typed `HydratedChunk`, which + retires the `best_distance` compatibility branch and the per-backend + `_distance_to_similarity` hooks the adapters had already replaced. Test doubles + construct the pipeline directly instead of subclassing the repository. No query + behavior changes. + ## v0.23.2 (2026-08-25) diff --git a/src/basic_memory/repository/postgres_search_repository.py b/src/basic_memory/repository/postgres_search_repository.py index 0cb17fcb0..06f4e57d8 100644 --- a/src/basic_memory/repository/postgres_search_repository.py +++ b/src/basic_memory/repository/postgres_search_repository.py @@ -490,15 +490,6 @@ async def _delete_stale_chunks( expected_deletions=expected_deletions, ) - @override - def _distance_to_similarity(self, distance: float) -> float: - """Convert pgvector cosine distance to cosine similarity. - - pgvector's <=> operator returns cosine distance in [0, 2], - where cos_distance = 1 - cos_similarity. - """ - return max(0.0, 1.0 - distance) - @override def _timestamp_now_expr(self) -> str: return "NOW()" diff --git a/src/basic_memory/repository/search_query.py b/src/basic_memory/repository/search_query.py index 9cd2cb2b4..59dd16182 100644 --- a/src/basic_memory/repository/search_query.py +++ b/src/basic_memory/repository/search_query.py @@ -32,6 +32,28 @@ class PreparedSearchQuery: retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS min_similarity: float | None = None + @property + def has_filters(self) -> bool: + """Whether any predicate beyond the text itself narrows the result set. + + Vector retrieval cannot evaluate these itself; when any is present it asks + the full-text pass which of its candidates the filters admit. + """ + return any( + ( + self.permalink, + self.permalink_match, + self.title, + self.note_types, + self.after_date, + self.search_item_types, + self.categories, + self.metadata_filters, + self.file_path_prefix, + self.temporal, + ) + ) + # Interrogative/function words contribute lexical noise when a strict # full-text query is relaxed: "when OR did OR a" matches loud wrong documents diff --git a/src/basic_memory/repository/search_reader.py b/src/basic_memory/repository/search_reader.py new file mode 100644 index 000000000..d5f925d54 --- /dev/null +++ b/src/basic_memory/repository/search_reader.py @@ -0,0 +1,1136 @@ +"""Scoped search retrieval shared by every route. + +``SearchReader`` runs one prepared query over one ``ProjectScope``: the engine's +full-text statement for FTS mode, and vector or hybrid retrieval through +``SemanticSearch`` when the semantic stack is present. Neither class owns indexing, +manifest writes, or table lifecycle. A repository builds a reader per call from its +current state, so a repository whose semantic stack was disabled at startup hands the +reader the matching capability set; a route reading several projects at once builds +one directly over a wider scope. +""" + +import time +from collections.abc import Sequence +from dataclasses import dataclass, replace +from typing import Any, assert_never + +import logfire +from loguru import logger +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from basic_memory import db +from basic_memory.repository.embedding_provider import EmbeddingProvider +from basic_memory.repository.rerank_provider import ( + RerankProvider, + build_rerank_document, + demote_tail_scores, + validate_rerank_scores, +) +from basic_memory.repository.search_filters import FtsBackend +from basic_memory.repository.search_index_row import SearchIndexKey, SearchIndexRow +from basic_memory.repository.search_query import PreparedSearchQuery +from basic_memory.repository.search_scope import ProjectScope +from basic_memory.repository.search_trace import ( + BelowThreshold, + FilteredOut, + HydrationDropKey, + HydrationDropped, + MissingSearchRow, + SearchTraceCollector, + build_fts_page_stage, + build_fusion_stage, + build_rerank_stage, + build_vector_stage, + classify_hydration_drops, + read_manifest_readiness, +) +from basic_memory.repository.semantic_errors import SemanticSearchDisabledError +from basic_memory.repository.semantic_vector_index import ( + SemanticVectorIndex, + VectorKey, + VectorMatch, +) +from basic_memory.schemas.search import SearchRetrievalMode + +# --- Retrieval constants --- + +# Adapters that share the authoritative SQL database. Everything else stores vectors +# outside it and needs the stale-hit overfetch below. +BUILT_IN_VECTOR_INDEX_NAMES = frozenset({"pgvector", "sqlite-vec"}) +VECTOR_FILTER_SCAN_LIMIT = 50000 +# The shared bind-parameter bound for any statement that carries a list of vector +# candidate keys. Both engines cap bind parameters (asyncpg at 32767), so every such +# list — manifest hydration and the filter intersection alike — is split at this size. +VECTOR_HYDRATION_BATCH_SIZE = 250 +# Over-fetch factor for the rerank candidate chunk pool: chunks collapse to unique +# (type, id) rows before reranking, so fetch several times reranker_candidates chunks +# to keep enough unique documents in the rerank window. +RERANK_POOL_CHUNK_FANOUT = 4 +FUSION_BONUS = 0.3 +FUSION_FORMULA_VERSION = "max+0.3*min/v1" +FTS_GATE_THRESHOLD = 0.0 +TOP_CHUNKS_PER_RESULT = 5 +SMALL_NOTE_CONTENT_LIMIT = 2000 + + +# The manifest conditions under which semantic retrieval will use a stored vector. +# Vector hydration admits exactly these rows, so anything failing them is invisible +# to search: a chunk left behind by an embedding-model or vector-index change, or one +# still pending. Readiness reporting must apply the same predicate — calling such a +# row "embedded" would report an index settled that retrieval cannot answer from, +# which is the class of lie #1414 exists to remove. +# Callers bind :vector_index and :embedding_model; the scope binds its own IDs. +def current_vector_manifest_predicate(scope: ProjectScope, params: dict[str, Any]) -> str: + """SQL admitting only manifest rows retrieval can answer from, within ``scope``.""" + return ( + f"{scope.predicate('project_id', params)} " + "AND vector_index = :vector_index " + "AND embedding_model = :embedding_model " + "AND embedding_status = 'ready'" + ) + + +def parse_chunk_key(chunk_key: str) -> SearchIndexKey: + """Parse a chunk key like ``observation:5:0`` into ``(type, search_index_id)``.""" + parts = chunk_key.split(":") + return parts[0], int(parts[1]) + + +def vector_eligible(query: PreparedSearchQuery) -> bool: + """Whether a query carries text to embed and no identity filter. + + Vector and hybrid retrieval score an embedding of the query text; a permalink or + title lookup has nothing to embed and answers exactly through the full-text path. + """ + text_value = (query.search_text or "").strip() + return ( + bool(text_value) + and text_value != "*" + and not query.permalink + and not query.permalink_match + and not query.title + ) + + +def rerank_document_text(row: SearchIndexRow, max_chars: int) -> str: + """Build the document text handed to the cross-encoder for one candidate. + + Prefer the matched chunk (the most relevant passage of a large note), falling + back to the stored snippet. + """ + body = row.matched_chunk_text or row.content_snippet or "" + return build_rerank_document(row.title, body, max_chars) + + +def demote_tail(tail: list[SearchIndexRow], floor: float) -> list[SearchIndexRow]: + """Rescore un-reranked tail rows at or below the floor, preserving their order. + + The reranked pool carries [0, 1] relevance scores while the tail still holds raw + retrieval scores on a different scale ([0, 1.3] for fused hybrid). Left as is, a + tail row could outrank a reranked row numerically. Positive floors put the tail + strictly below the pool; a zero floor yields zeroes because no smaller score + exists in the public [0, 1] range. The returned pool-plus-tail sequence, rather + than a later score-only sort, owns that tie-breaking invariant. + """ + return [ + replace(row, score=score) for row, score in zip(tail, demote_tail_scores(floor, len(tail))) + ] + + +# --- Retrieval capabilities --- + + +@dataclass(frozen=True, slots=True) +class VectorRetrieval: + """The live semantic stack vector retrieval reads from. + + Present only when semantic search is enabled, a provider is configured, and the + adapter is bound; absence means the reader answers full-text queries only. + """ + + index: SemanticVectorIndex + index_name: str + embedding_provider: EmbeddingProvider + # The persisted embedding identity manifest rows are keyed by. + embedding_model: str + vector_k: int + min_similarity: float + + @property + def external(self) -> bool: + """Whether vectors live outside the SQL database that owns the manifest.""" + return self.index_name not in BUILT_IN_VECTOR_INDEX_NAMES + + +@dataclass(frozen=True, slots=True) +class Reranking: + """A configured cross-encoder and the fixed prefix it rescores.""" + + provider: RerankProvider + candidates: int + max_document_chars: int + + +@dataclass(frozen=True, slots=True) +class HydratedChunk: + """One adapter match the ready manifest confirmed retrieval may serve.""" + + entity_id: int + chunk_key: str + chunk_text: str + similarity: float + + +# --- Vector and hybrid retrieval --- + + +class SemanticSearch: + """Vector and hybrid retrieval over one scope. + + Constructed only when the semantic stack is present, so every stage reads its + adapter, provider, and thresholds directly instead of re-checking availability. + """ + + def __init__( + self, + session_maker: async_sessionmaker[AsyncSession], + scope: ProjectScope, + fts: FtsBackend, + vector: VectorRetrieval, + rerank: Reranking | None = None, + ) -> None: + self.session_maker = session_maker + self.scope = scope + self.fts = fts + self.vector = vector + self.rerank = rerank + + # --- Candidate window sizing --- + + def _active_rerank(self, query_text: str) -> Reranking | None: + """The reranker this query runs: one is configured and there is text to score.""" + return self.rerank if query_text else None + + def _rerank_candidate_limit(self, rerank: Reranking) -> int: + """Return the fixed chunk window that owns reranker-prefix membership.""" + return max(self.vector.vector_k, rerank.candidates * RERANK_POOL_CHUNK_FANOUT) + + def _candidate_limit(self, limit: int, offset: int, query_text: str) -> int: + """Size the retrieval candidate *chunk* pool for vector/hybrid search. + + ``candidate_limit`` bounds vector chunks, but many chunks of one large note + collapse to a single ``(type, id)`` row before reranking, so a chunk count does + not equal a unique-document count. When reranking is active we over-fetch by + ``RERANK_POOL_CHUNK_FANOUT`` so a few multi-chunk notes can't starve the rerank + window below ``reranker_candidates`` unique rows. This is best-effort headroom, + not a hard guarantee — a single note dominating the entire nearest-neighbour set + can still yield fewer unique rows (a pathological corpus shape). + """ + rerank = self._active_rerank(query_text) + if rerank is None: + return max(self.vector.vector_k, (limit + offset) * 10) + # Trigger: the requested window extends beyond the fixed reranked prefix. + # Why: a bounded prefix alone can under-fill large pages and hide the + # semantic pagination probe even when more matches exist. + # Outcome: keep prefix membership fixed while adding chunk headroom only + # for the untouched tail that this request must return. + tail_size = max(0, limit + offset - rerank.candidates) + return self._rerank_candidate_limit(rerank) + tail_size * 10 + + # --- Vector nearest neighbours through the ready manifest --- + + @logfire.instrument("search.vector_query", extract_args=False) + async def _run_vector_query( + self, + session: AsyncSession, + query_embedding: list[float], + candidate_limit: int, + *, + trace: SearchTraceCollector | None = None, + ) -> list[HydratedChunk]: + """Query the configured adapter and hydrate only live, ready manifest rows.""" + if trace is not None: + trace.vector = build_vector_stage( + candidate_limit=candidate_limit, + adapter_match_count=0, + hydrated_count=0, + ) + if candidate_limit <= 0: + return [] + + if not self.vector.external: + matches = await self.vector.index.search(query_embedding, limit=candidate_limit) + if trace is not None: + trace.readiness = await read_manifest_readiness( + session, + self.scope, + self.vector.index_name, + self.vector.embedding_model, + ) + return await self._hydrate_vector_matches(session, matches, trace=trace) + + scan_limit = min(candidate_limit, VECTOR_FILTER_SCAN_LIMIT) + while True: + matches = await self.vector.index.search(query_embedding, limit=scan_limit) + if trace is not None and trace.readiness is None: + trace.readiness = await read_manifest_readiness( + session, + self.scope, + self.vector.index_name, + self.vector.embedding_model, + ) + hydrated = await self._hydrate_vector_matches(session, matches, trace=trace) + if ( + len(hydrated) >= candidate_limit + or len(matches) < scan_limit + or scan_limit >= VECTOR_FILTER_SCAN_LIMIT + ): + returned = hydrated[:candidate_limit] + # Trigger: the expanded stale-hit rescan hydrated more chunks than the + # candidate window the search consumes. + # Why: chunks beyond the window never enter thresholding, fusion, or + # reranking — tracing them would invent candidates this execution + # never considered. + # Outcome: the traced stage is trimmed to the returned window. + if trace is not None and trace.vector is not None and len(hydrated) > len(returned): + # Two owners can share one parseable chunk_key (manifest uniqueness + # includes entity_id), so window membership matches by owner too. + returned_chunk_keys = {(chunk.entity_id, chunk.chunk_key) for chunk in returned} + trimmed: dict[SearchIndexKey, list[tuple[str, float, int | None]]] = {} + for chunk_match in trace.vector.chunk_matches: + if (chunk_match.entity_id, chunk_match.chunk_key) in returned_chunk_keys: + trimmed.setdefault(chunk_match.key, []).append( + ( + chunk_match.chunk_key, + chunk_match.similarity, + chunk_match.entity_id, + ) + ) + # hydrated_count keeps full-scan scope so the vector stage's + # dropped count matches its hydration-drop list; the flattener + # reports the window truncation as its own candidate_window stage. + trace.vector = build_vector_stage( + previous=trace.vector, + chunk_matches=trimmed, + ) + return returned + + # Trigger: stale, pending, or wrong-model adapter hits consumed the + # requested top-k before manifest hydration. + # Why: returning early lets stale extension data crowd every live + # result out of an otherwise valid semantic search. + # Outcome: retry from the same ranked prefix with bounded geometric + # overfetch until enough live rows survive or the adapter is exhausted. + scan_limit = min(scan_limit * 2, VECTOR_FILTER_SCAN_LIMIT) + + @logfire.instrument("search.vector_manifest_hydration", extract_args=False) + async def _hydrate_vector_matches( + self, + session: AsyncSession, + matches: list[VectorMatch], + *, + trace: SearchTraceCollector | None = None, + ) -> list[HydratedChunk]: + """Resolve adapter matches through the authoritative ready manifest.""" + if not matches: + return [] + + chunks_by_key: dict[VectorKey, str] = {} + for batch_start in range(0, len(matches), VECTOR_HYDRATION_BATCH_SIZE): + batch = matches[batch_start : batch_start + VECTOR_HYDRATION_BATCH_SIZE] + params: dict[str, Any] = { + "vector_index": self.vector.index_name, + "embedding_model": self.vector.embedding_model, + } + manifest_predicate = current_vector_manifest_predicate(self.scope, params) + predicates: list[str] = [] + for index, match in enumerate(batch): + params[f"entity_id_{index}"] = match.key.entity_id + params[f"chunk_key_{index}"] = match.key.chunk_key + predicates.append( + f"(entity_id = :entity_id_{index} AND chunk_key = :chunk_key_{index})" + ) + + # Constraint: adapters may return thousands of candidates for deep pages. + # PostgreSQL and SQLite both cap bind parameters, so hydrate in fixed-size + # batches while retaining the adapter's original ranking in the final list. + result = await session.execute( + text( + "SELECT entity_id, chunk_key, chunk_text FROM search_vector_chunks " + "WHERE " + manifest_predicate + " " + "AND (" + " OR ".join(predicates) + ")" + ), + params, + ) + chunks_by_key.update( + { + VectorKey( + entity_id=int(row["entity_id"]), + chunk_key=str(row["chunk_key"]), + ): str(row["chunk_text"]) + for row in result.mappings().all() + } + ) + hydrated = [ + HydratedChunk( + entity_id=match.key.entity_id, + chunk_key=match.key.chunk_key, + chunk_text=chunks_by_key[match.key], + similarity=match.similarity, + ) + for match in matches + if match.key in chunks_by_key + ] + if trace is not None: + dropped_keys = [ + HydrationDropKey( + entity_id=match.key.entity_id, + chunk_key=match.key.chunk_key, + similarity=match.similarity, + configured_index=self.vector.index_name, + configured_model=self.vector.embedding_model, + ) + for match in matches + if match.key not in chunks_by_key + ] + drops = await classify_hydration_drops(session, self.scope, dropped_keys) + chunk_matches: dict[SearchIndexKey, list[tuple[str, float, int | None]]] = {} + malformed_drops: list[HydrationDropped] = [] + for chunk in hydrated: + try: + key = parse_chunk_key(chunk.chunk_key) + except (ValueError, IndexError): + # A hydrated chunk with an unparseable key silently vanishes from + # retrieval; the trace must name it or the stage counts lie. + malformed_drops.append( + HydrationDropped( + entity_id=chunk.entity_id, + chunk_key=chunk.chunk_key, + similarity=chunk.similarity, + reason="malformed_key", + stored_model=None, + stored_index=None, + ) + ) + continue + chunk_matches.setdefault(key, []).append( + (chunk.chunk_key, chunk.similarity, chunk.entity_id) + ) + trace.vector = build_vector_stage( + previous=trace.vector, + adapter_match_count=len(matches), + # Malformed keys are dropped, not served — counting them as output + # would contradict the malformed_key rejection listed alongside. + hydrated_count=len(hydrated) - len(malformed_drops), + drops=(*drops, *malformed_drops), + chunk_matches=chunk_matches, + ) + return hydrated + + # --- Candidate rows and structured filters --- + + @logfire.instrument("search.fetch_candidate_rows", extract_args=False) + async def _fetch_search_index_rows_by_ids( + self, row_ids: list[int] + ) -> dict[SearchIndexKey, SearchIndexRow]: + """Fetch search_index rows by id, keyed by (type, id) to disambiguate types. + + A bare id can match one row per type (independent id sequences), so the + result must carry every matching row rather than letting one clobber another. + """ + if not row_ids: + return {} + placeholders = ",".join(f":id_{idx}" for idx in range(len(row_ids))) + params: dict[str, Any] = {f"id_{idx}": rid for idx, rid in enumerate(row_ids)} + scope_predicate = self.scope.predicate("project_id", params) + sql = f""" + SELECT + project_id, id, title, permalink, file_path, type, metadata, + from_id, to_id, relation_type, entity_id, content_snippet, + category, created_at, updated_at, 0 as score + FROM search_index + WHERE {scope_predicate} + AND id IN ({placeholders}) + """ + result: dict[SearchIndexKey, SearchIndexRow] = {} + async with db.scoped_session(self.session_maker) as session: + row_result = await session.execute(text(sql), params) + for row in row_result.fetchall(): + search_row = SearchIndexRow.from_mapping(row._asdict()) + result[(search_row.type, search_row.id)] = search_row + return result + + @logfire.instrument("search.filter_candidates", extract_args=False) + async def _filter_candidate_keys( + self, + candidate_keys: Sequence[SearchIndexKey], + query: PreparedSearchQuery, + ) -> set[SearchIndexKey]: + """Return which of ``candidate_keys`` the query's structured filters admit. + + Vector retrieval scores embeddings and cannot evaluate a structured filter, so + the surviving candidates are decided by an FTS-mode pass carrying every filter. + Asking that pass for a *page of the filter's whole match set* and intersecting + client-side silently lost any candidate that sorted past the page (#1431); asking + it about the candidates themselves cannot, because the answer is bounded by the + question. + + The candidate list is split at the shared bind-parameter bound, so a deep page + whose candidate pool runs to thousands of rows costs a few small indexed lookups + instead of one unbounded scan. + """ + filter_query = replace(query, search_text=None, retrieval_mode=SearchRetrievalMode.FTS) + allowed_keys: set[SearchIndexKey] = set() + for batch_start in range(0, len(candidate_keys), VECTOR_HYDRATION_BATCH_SIZE): + batch = candidate_keys[batch_start : batch_start + VECTOR_HYDRATION_BATCH_SIZE] + filtered_rows = await self.fts.search( + self.scope, + filter_query, + # The restriction, not this limit, is what bounds the result: one row per + # requested key, since (id, type, project_id) identifies a search row. + limit=len(batch), + offset=0, + candidate_keys=batch, + ) + allowed_keys.update((row.type, row.id) for row in filtered_rows if row.id is not None) + return allowed_keys + + # --- Reranking --- + + async def _rerank_and_paginate( + self, + query_text: str, + rows: list[SearchIndexRow], + *, + offset: int, + limit: int, + stable_rows: list[SearchIndexRow] | None = None, + trace: SearchTraceCollector | None = None, + ) -> list[SearchIndexRow]: + """Rerank the top candidates, then return the requested ``[offset:offset+limit]`` page. + + Trigger: a reranker is configured and there is a real query. + Why: bi-encoder/FTS ranking lands the gold document in the top-N but often + just below the top-k cutoff (#950); a cross-encoder that reads query and + document together recovers those near-misses. + Outcome: the first ``reranker_candidates`` rows are reordered by reranker + relevance (which replaces ``score``); the requested page is sliced from the + reordered list. + + Every non-empty page rescores the same fixed prefix before slicing so the + untouched tail can be demoted onto the reranker's public ``[0, 1]`` scale. + """ + page_end = offset + limit + rerank = self._active_rerank(query_text) + if rerank is None: + return rows[offset:page_end] + + # Trigger: pagination needs more rows than the fixed rerank retrieval window. + # Why: an expanded retrieval may introduce or strengthen raw candidates, but + # letting them replace the original prefix causes duplicates and skips. + # Outcome: the fixed window owns prefix membership; the expanded result only + # supplies new, de-duplicated tail rows. + pool_source = stable_rows if stable_rows is not None else rows + pool = pool_source[: rerank.candidates] + pool_keys = {(row.type, row.id) for row in pool} + tail = [row for row in rows if (row.type, row.id) not in pool_keys] + ordered_rows = pool + tail + + # Skip only when there is no prefix to calibrate or the requested page is + # empty. Even a singleton prefix or a wholly-tail page needs the prefix's + # relevance floor so raw hybrid scores cannot leak into cross-project sorting. + if not pool or offset >= len(ordered_rows): + return ordered_rows[offset:page_end] + + pre_rerank_scores = None + if trace is not None: + pre_rerank_scores = {(row.type, row.id): row.score or 0.0 for row in ordered_rows} + documents = [rerank_document_text(row, rerank.max_document_chars) for row in pool] + # A transient provider failure must surface instead of switching this page + # back to retrieval order. A prior page may already have returned reranked + # order, so degrading here can duplicate one result and omit another. + rerank_start = time.perf_counter() if trace is not None else None + with logfire.span( + "search.rerank", + candidate_count=len(pool), + document_chars=sum(map(len, documents)), + ): + scores = validate_rerank_scores( + await rerank.provider.rerank(query_text, documents), + len(pool), + ) + + order = sorted(range(len(pool)), key=lambda i: scores[i], reverse=True) + reranked = [replace(pool[i], score=scores[i]) for i in order] + logger.debug( + "Reranked candidates: pool={pool} model={model}", + pool=len(pool), + model=rerank.provider.model_name, + ) + tail_floor = reranked[-1].score or 0.0 + demoted_tail = demote_tail(tail, floor=tail_floor) + reranked_rows = reranked + demoted_tail + if trace is not None: + assert pre_rerank_scores is not None and rerank_start is not None + trace.rerank = build_rerank_stage( + provider_model=rerank.provider.model_name, + reranker_candidates=rerank.candidates, + pre_rerank_scores=pre_rerank_scores, + pool_keys=[(row.type, row.id) for row in pool], + rerank_scores={ + (pool[index].type, pool[index].id): score for index, score in enumerate(scores) + }, + post_rerank_rows=[((row.type, row.id), row.score or 0.0) for row in reranked_rows], + demoted_scores={(row.type, row.id): row.score or 0.0 for row in demoted_tail}, + tail_floor=tail_floor, + stable_pool_refetched=trace.stable_pool_refetched, + rerank_ms=(time.perf_counter() - rerank_start) * 1000, + ) + return reranked_rows[offset:page_end] + + # --- Vector-only retrieval --- + + async def vector_only( + self, + query: PreparedSearchQuery, + *, + limit: int, + offset: int, + candidate_limit: int | None = None, + apply_rerank: bool = True, + emit_observability_log: bool = True, + trace: SearchTraceCollector | None = None, + ) -> list[SearchIndexRow]: + """Run vector-only search returning chunk-level results. + + Returns individual search_index rows (entities, observations, relations) + ranked by vector similarity. Each observation or relation is a first-class + result, not collapsed into its parent entity. + + ``candidate_limit`` is supplied only by a composed retrieval stage that + already sized the shared candidate pool. + """ + query_text = (query.search_text or "").strip() + if candidate_limit is None: + candidate_limit = self._candidate_limit(limit, offset, query_text) + query_start = time.perf_counter() + embed_start = time.perf_counter() + with logfire.span("search.embed_query", query_chars=len(query_text)): + query_embedding = await self.vector.embedding_provider.embed_query(query_text) + embed_ms = (time.perf_counter() - embed_start) * 1000 + vector_query_start = time.perf_counter() + + # Constraint: vector adapters may open their own session, while the SQLite + # test/runtime pool can contain only one connection. A plain AsyncSession + # defers checkout until hydration runs after adapter search has released it. + async with self.session_maker() as session: + vector_rows = await self._run_vector_query( + session, + query_embedding, + candidate_limit, + trace=trace, + ) + vector_query_ms = (time.perf_counter() - vector_query_start) * 1000 + vector_row_count = len(vector_rows) + hydrate_ms = 0.0 + + if trace is not None: + trace.vector = build_vector_stage( + previous=trace.vector, + effective_min_similarity=( + query.min_similarity + if query.min_similarity is not None + else self.vector.min_similarity + ), + min_similarity_source=("query" if query.min_similarity is not None else "config"), + embed_ms=embed_ms, + vector_query_ms=vector_query_ms, + ) + + def _log_vector_summary() -> None: + if not emit_observability_log: + return + + total_ms = (time.perf_counter() - query_start) * 1000 + if total_ms > 2000: + logger.warning( + "[SEMANTIC_SLOW_QUERY] Semantic query timing: scope={scope} " + "retrieval_mode={retrieval_mode} query_length={query_length} " + "candidate_limit={candidate_limit} vector_row_count={vector_row_count} " + "embed_ms={embed_ms:.2f} vector_query_ms={vector_query_ms:.2f} " + "hydrate_ms={hydrate_ms:.2f} total_ms={total_ms:.2f}", + scope=self.scope.project_ids, + retrieval_mode="vector", + query_length=len(query_text), + candidate_limit=candidate_limit, + vector_row_count=vector_row_count, + embed_ms=embed_ms, + vector_query_ms=vector_query_ms, + hydrate_ms=hydrate_ms, + total_ms=total_ms, + ) + + if not vector_rows: + _log_vector_summary() + return [] + + hydrate_start = time.perf_counter() + # Build per-search_index_row similarity scores from chunk-level results. + # Each chunk_key encodes the search_index row type and id; keep both as the + # key because different row types can share the same numeric id (#982). + # Track the best similarity per row (for ranking) and all chunks (for context). + similarity_by_si_key: dict[SearchIndexKey, float] = {} + chunks_by_si_key: dict[SearchIndexKey, list[tuple[float, str]]] = {} + for chunk in vector_rows: + try: + si_key = parse_chunk_key(chunk.chunk_key) + except (ValueError, IndexError): + # A chunk without a parseable key names no search row to rank. + continue + current = similarity_by_si_key.get(si_key) + if current is None or chunk.similarity > current: + similarity_by_si_key[si_key] = chunk.similarity + chunks_by_si_key.setdefault(si_key, []).append((chunk.similarity, chunk.chunk_text)) + + if not similarity_by_si_key: + hydrate_ms = (time.perf_counter() - hydrate_start) * 1000 + _log_vector_summary() + return [] + + # Filter out results below the minimum similarity threshold. + # Per-query min_similarity overrides the configured default. + effective_min_similarity = ( + query.min_similarity if query.min_similarity is not None else self.vector.min_similarity + ) + if effective_min_similarity > 0.0: + if trace is not None: + threshold_rejections = tuple( + BelowThreshold(key=key, similarity=value, threshold=effective_min_similarity) + for key, value in similarity_by_si_key.items() + if value < effective_min_similarity + ) + trace.vector = build_vector_stage( + previous=trace.vector, + threshold_rejections=threshold_rejections, + ) + similarity_by_si_key = { + k: v for k, v in similarity_by_si_key.items() if v >= effective_min_similarity + } + if not similarity_by_si_key: + hydrate_ms = (time.perf_counter() - hydrate_start) * 1000 + _log_vector_summary() + return [] + + # Fetch the actual search_index rows. Colliding (type, id) keys share one + # bare id, so deduplicate while preserving first-seen order. + si_ids = list(dict.fromkeys(si_id for _, si_id in similarity_by_si_key)) + search_index_rows = await self._fetch_search_index_rows_by_ids(si_ids) + if trace is not None: + trace.vector = build_vector_stage( + previous=trace.vector, + missing_search_rows=tuple( + MissingSearchRow(key=key) + for key in similarity_by_si_key + if key not in search_index_rows + ), + ) + + if query.has_filters: + allowed_keys = await self._filter_candidate_keys(list(search_index_rows), query) + if trace is not None: + trace.vector = build_vector_stage( + previous=trace.vector, + filter_rejections=tuple( + FilteredOut(key=key) for key in search_index_rows if key not in allowed_keys + ), + ) + search_index_rows = {k: v for k, v in search_index_rows.items() if k in allowed_keys} + + ranked_rows: list[SearchIndexRow] = [] + for si_key, similarity in similarity_by_si_key.items(): + row = search_index_rows.get(si_key) + if row is None: + continue + + # Small notes: return full content so the answer is always present. + # Large notes: return top-N most relevant chunks for richer context. + content_snippet = row.content_snippet or "" + if content_snippet and len(content_snippet) <= SMALL_NOTE_CONTENT_LIMIT: + matched_chunk_text = content_snippet + else: + si_chunks = chunks_by_si_key.get(si_key, []) + si_chunks.sort(key=lambda c: c[0], reverse=True) + top_texts = [chunk_text for _, chunk_text in si_chunks[:TOP_CHUNKS_PER_RESULT]] + matched_chunk_text = "\n---\n".join(top_texts) if top_texts else None + + ranked_rows.append( + replace( + row, + score=similarity, + matched_chunk_text=matched_chunk_text, + ) + ) + + ranked_rows.sort(key=lambda item: item.score or 0.0, reverse=True) + hydrate_ms = (time.perf_counter() - hydrate_start) * 1000 + # Rerank over the wide candidate pool, then slice to the page. Suppressed when + # hybrid calls this internally (apply_rerank=False): hybrid reranks its own + # fused result, and _rerank_and_paginate is a plain slice without a reranker. + if apply_rerank: + stable_rows = ranked_rows + rerank = self._active_rerank(query_text) + if rerank is not None: + stable_candidate_limit = self._rerank_candidate_limit(rerank) + if candidate_limit > stable_candidate_limit: + if trace is not None: + trace.stable_pool_refetched = True + stable_rows = await self.vector_only( + query, + limit=stable_candidate_limit, + offset=0, + candidate_limit=stable_candidate_limit, + apply_rerank=False, + emit_observability_log=False, + trace=None, + ) + output = await self._rerank_and_paginate( + query_text, + ranked_rows, + offset=offset, + limit=limit, + stable_rows=stable_rows, + trace=trace, + ) + else: + output = ranked_rows[offset : offset + limit] + # Vector latency owns the optional rerank stage too. Logging before the + # awaited provider call hides the feature's dominant cost and can suppress + # the slow-query warning entirely. + _log_vector_summary() + return output + + # --- Hybrid score-based fusion --- + + async def hybrid( + self, + query: PreparedSearchQuery, + *, + limit: int, + offset: int, + candidate_limit: int | None = None, + apply_rerank: bool = True, + emit_observability_log: bool = True, + trace: SearchTraceCollector | None = None, + ) -> list[SearchIndexRow]: + """Fuse FTS and vector results using score-based fusion. + + Uses the search_index (type, id) pair as the fusion key. The formula + ``max(vec, fts) + FUSION_BONUS * min(vec, fts)`` preserves + the dominant signal and rewards dual-source agreement. + """ + query_text = (query.search_text or "").strip() + rerank = self._active_rerank(query_text) + query_start = time.perf_counter() + if candidate_limit is None: + candidate_limit = self._candidate_limit(limit, offset, query_text) + fts_start = time.perf_counter() + # allow_relaxed: question-form queries rarely AND-match, and a dead FTS + # branch silently degrades hybrid to vector-only ranking. Fusion plus + # bm25 keep relaxed lexical candidates from dominating precision. + with logfire.span("search.fts", candidate_limit=candidate_limit) as fts_span: + fts_results = await self.fts.search( + self.scope, + replace(query, retrieval_mode=SearchRetrievalMode.FTS), + limit=candidate_limit, + offset=0, + allow_relaxed=True, + trace=trace, + ) + fts_span.set_attribute("result_count", len(fts_results)) + fts_ms = (time.perf_counter() - fts_start) * 1000 + vector_start = time.perf_counter() + vector_results = await self.vector_only( + query, + limit=candidate_limit, + offset=0, + # Trigger: reranking owns a bounded candidate window shared by both legs. + # Why: the disabled path historically expands the vector leg again to + # preserve recall when many vector chunks collapse into a few search rows. + # Outcome: avoid double expansion only when reranking is actually active. + candidate_limit=candidate_limit if rerank is not None else None, + apply_rerank=False, + emit_observability_log=False, + trace=trace, + ) + vector_ms = (time.perf_counter() - vector_start) * 1000 + # Trigger: with reranking disabled the vector leg expands internally and can + # hydrate more rows than the fusion window it returns. + # Why: rows cut here never fuse — left in the trace they would surface as + # candidates with no rejection and no fused rank, which the response labels + # "returned". Rows with a recorded rejection keep their chunk evidence. + # Outcome: the trace keeps rows handed to fusion (or explicitly rejected); + # the cut shows up as served-chunk shrinkage in the candidate_window stage. + if trace is not None and trace.vector is not None: + kept_row_keys = {(row.type, row.id) for row in vector_results} + kept_row_keys.update( + rejection.key + for rejection_group in ( + trace.vector.threshold_rejections, + trace.vector.filter_rejections, + trace.vector.missing_search_rows, + ) + for rejection in rejection_group + ) + if any(match.key not in kept_row_keys for match in trace.vector.chunk_matches): + fused_chunks: dict[SearchIndexKey, list[tuple[str, float, int | None]]] = {} + for chunk_match in trace.vector.chunk_matches: + if chunk_match.key in kept_row_keys: + fused_chunks.setdefault(chunk_match.key, []).append( + (chunk_match.chunk_key, chunk_match.similarity, chunk_match.entity_id) + ) + trace.vector = build_vector_stage( + previous=trace.vector, + chunk_matches=fused_chunks, + ) + fusion_start = time.perf_counter() + + with logfire.span( + "search.fusion", fts_count=len(fts_results), vector_count=len(vector_results) + ) as fusion_span: + # --- Score-based fusion keyed on (type, id) --- + # A bare row id collides across row types (independent id sequences), so + # fusion must key on (type, id) or distinct rows would merge (#982). + # FTS scores are normalized to [0, 1] (BM25 is unbounded). + # Vector scores are used raw: the adapters already calibrate them to [0, 1]. + rows_by_key: dict[SearchIndexKey, SearchIndexRow] = {} + + # Normalize FTS scores to [0, 1] — handles both SQLite (negative bm25) + # and Postgres (positive ts_rank) by using absolute values + fts_abs = [abs(row.score or 0.0) for row in fts_results] + fts_max = max(fts_abs) if fts_abs else 1.0 + + fts_scores: dict[SearchIndexKey, float] = {} + fts_ranks: dict[SearchIndexKey, int] = {} + for rank, row in enumerate(fts_results): + if row.id is None: + continue + row_key = (row.type, row.id) + norm = abs(row.score or 0.0) / fts_max if fts_max > 0 else 0.0 + # Gate: FTS scores below threshold contribute zero + if norm < FTS_GATE_THRESHOLD: + norm = 0.0 + fts_scores[row_key] = norm + fts_ranks.setdefault(row_key, rank) + rows_by_key[row_key] = row + + if trace is not None: + relaxed_fallback_used = ( + trace.fts.relaxed_fallback_used if trace.fts is not None else False + ) + trace.fts = build_fts_page_stage( + [((row.type, row.id), row.score or 0.0) for row in fts_results], + normalized_scores=fts_scores, + entity_ids={(row.type, row.id): row.entity_id for row in fts_results}, + fts_max_abs=fts_max, + relaxed_fallback_used=relaxed_fallback_used, + fts_ms=fts_ms, + ) + + vec_scores: dict[SearchIndexKey, float] = {} + vec_ranks: dict[SearchIndexKey, int] = {} + for rank, row in enumerate(vector_results): + if row.id is None: + continue + row_key = (row.type, row.id) + # Trigger: no re-normalization by vec_max + # Why: vector similarity is already calibrated [0, 1]; re-normalizing + # inflates weak matches when the entire result set is mediocre + vec_scores[row_key] = row.score or 0.0 + vec_ranks.setdefault(row_key, rank) + rows_by_key[row_key] = row + + # Fuse: max(v, f) + FUSION_BONUS * min(v, f) + # Preserves the dominant signal; bonus rewards dual-source agreement. + # Output range: [0, 1.3] for dual-source, [0, 1.0] for single-source. + fused_scores: dict[SearchIndexKey, float] = {} + for row_key in fts_scores.keys() | vec_scores.keys(): + v = vec_scores.get(row_key, 0.0) + f = fts_scores.get(row_key, 0.0) + fused_scores[row_key] = max(v, f) + FUSION_BONUS * min(v, f) + + ranked = sorted(fused_scores.items(), key=lambda item: item[1], reverse=True) + fusion_span.set_attribute("result_count", len(ranked)) + fusion_ms = (time.perf_counter() - fusion_start) * 1000 + if trace is not None: + trace.fusion = build_fusion_stage( + formula_version=FUSION_FORMULA_VERSION, + bonus=FUSION_BONUS, + fts_scores=fts_scores, + fts_ranks=fts_ranks, + vector_scores=vec_scores, + vector_ranks=vec_ranks, + ranked_scores=ranked, + fusion_ms=fusion_ms, + ) + + def _materialize(entry: tuple[SearchIndexKey, float]) -> SearchIndexRow: + row_key, fused_score = entry + row = rows_by_key[row_key] + # FTS-only hits use the bounded content preview and its truncation metadata. + # Copying the full note into matched_chunk bypasses that response bound. + return replace(row, score=fused_score) + + # Rerank the top fused candidates before paginating. When reranking is active + # we materialize the whole candidate list (cheap next to a cross-encoder call) + # and hand it to the shared paginate helper; the disabled path stays cheap by + # materializing only the requested page. + if apply_rerank and rerank is not None: + candidates = [_materialize(entry) for entry in ranked] + stable_candidates = candidates + stable_candidate_limit = self._rerank_candidate_limit(rerank) + if candidate_limit > stable_candidate_limit: + if trace is not None: + trace.stable_pool_refetched = True + stable_candidates = await self.hybrid( + query, + limit=stable_candidate_limit, + offset=0, + candidate_limit=stable_candidate_limit, + apply_rerank=False, + emit_observability_log=False, + trace=None, + ) + stable_keys = {(row.type, row.id) for row in stable_candidates} + expanded_tail = [entry for entry in ranked if entry[0] not in stable_keys] + + # Trigger: deeper pages expand the FTS/vector retrieval windows. + # Why: score fusion can strengthen an existing row when its second + # signal appears later, moving it across a page already returned. + # Outcome: freeze the fixed fused universe, then order newly admitted + # rows by their earliest source rank. That rank cannot improve after a + # row first appears, so each larger window only appends to the tail. + expanded_tail.sort( + key=lambda entry: ( + min( + fts_ranks.get(entry[0], candidate_limit), + vec_ranks.get(entry[0], candidate_limit), + ), + entry[0], + ) + ) + candidates = stable_candidates + [_materialize(entry) for entry in expanded_tail] + output = await self._rerank_and_paginate( + query_text, + candidates, + offset=offset, + limit=limit, + stable_rows=stable_candidates, + trace=trace, + ) + else: + output = [_materialize(entry) for entry in ranked[offset : offset + limit]] + total_ms = (time.perf_counter() - query_start) * 1000 + if emit_observability_log and total_ms > 2500: + logger.warning( + "[SEMANTIC_SLOW_QUERY] Semantic query timing: scope={scope} " + "retrieval_mode={retrieval_mode} query_length={query_length} " + "candidate_limit={candidate_limit} fts_count={fts_count} " + "vector_count={vector_count} fts_ms={fts_ms:.2f} vector_ms={vector_ms:.2f} " + "fusion_ms={fusion_ms:.2f} total_ms={total_ms:.2f}", + scope=self.scope.project_ids, + retrieval_mode="hybrid", + query_length=len(query_text), + candidate_limit=candidate_limit, + fts_count=len(fts_results), + vector_count=len(vector_results), + fts_ms=fts_ms, + vector_ms=vector_ms, + fusion_ms=fusion_ms, + total_ms=total_ms, + ) + return output + + +# --- The reader --- + + +class SearchReader: + """Run one prepared query over one scope, whichever retrieval mode it asks for.""" + + def __init__( + self, + scope: ProjectScope, + fts: FtsBackend, + semantic: SemanticSearch | None = None, + ) -> None: + self.scope = scope + self.fts = fts + self.semantic = semantic + + def _semantic(self) -> SemanticSearch: + if self.semantic is None: + raise SemanticSearchDisabledError( + "Semantic search is disabled. Set BASIC_MEMORY_SEMANTIC_SEARCH_ENABLED=true." + ) + return self.semantic + + async def search( + self, + query: PreparedSearchQuery, + *, + limit: int, + offset: int, + allow_relaxed: bool = False, + session: AsyncSession | None = None, + candidate_keys: Sequence[SearchIndexKey] | None = None, + trace: SearchTraceCollector | None = None, + ) -> list[SearchIndexRow]: + """Search the scope in the query's retrieval mode. + + ``candidate_keys`` restricts full-text results to those ``(type, id)`` search + rows. ``None`` searches the whole scope; an empty sequence matches nothing. + Vector and hybrid retrieval use that restriction to ask which of a known + candidate set a filter admits instead of paging the filter's whole match set + (#1431). + + ``allow_relaxed=True`` retries a zero-result strict multi-word query with + OR-joined content terms. Only the hybrid path opts in: its FTS branch otherwise + contributes nothing for question-form queries. + """ + match query.retrieval_mode: + case SearchRetrievalMode.FTS: + return await self.fts.search( + self.scope, + query, + limit=limit, + offset=offset, + allow_relaxed=allow_relaxed, + session=session, + candidate_keys=candidate_keys, + trace=trace, + ) + case SearchRetrievalMode.VECTOR: + if not vector_eligible(query): + raise ValueError( + "Vector retrieval requires a non-empty text query and does not support " + "title/permalink-only searches." + ) + return await self._semantic().vector_only( + query, limit=limit, offset=offset, trace=trace + ) + case SearchRetrievalMode.HYBRID: + if not vector_eligible(query): + raise ValueError( + "Hybrid retrieval requires a non-empty text query and does not support " + "title/permalink-only searches." + ) + return await self._semantic().hybrid(query, limit=limit, offset=offset, trace=trace) + case _: # pragma: no cover + assert_never(query.retrieval_mode) + + async def count(self, query: PreparedSearchQuery, *, allow_relaxed: bool = False) -> int: + """Count full-text matches with the same filters as ``search``.""" + if query.retrieval_mode != SearchRetrievalMode.FTS: + raise ValueError("Exact counts are only supported for full-text search retrieval.") + return await self.fts.count(self.scope, query, allow_relaxed=allow_relaxed) diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 31f97096a..15e2dd506 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -5,7 +5,7 @@ from abc import ABC, abstractmethod from collections.abc import Iterable, Mapping, Sequence from contextlib import asynccontextmanager -from dataclasses import dataclass, replace +from dataclasses import dataclass from datetime import datetime, timezone from typing import Any, Callable, Dict, List, Literal, Optional, cast @@ -26,31 +26,22 @@ from basic_memory.repository.embedding_provider_factory import ( configured_embedding_provider_identity, ) -from basic_memory.repository.rerank_provider import ( - RerankProvider, - build_rerank_document, - demote_tail_scores, - validate_rerank_scores, -) +from basic_memory.repository.rerank_provider import RerankProvider from basic_memory.repository.search_filters import FtsBackend from basic_memory.repository.search_index_row import SearchIndexKey, SearchIndexRow from basic_memory.repository.search_query import PreparedSearchQuery +from basic_memory.repository.search_reader import ( + BUILT_IN_VECTOR_INDEX_NAMES, + VECTOR_HYDRATION_BATCH_SIZE, + Reranking, + SearchReader, + SemanticSearch, + VectorRetrieval, + vector_eligible, +) from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.script_ngrams import build_script_ngrams -from basic_memory.repository.search_trace import ( - BelowThreshold, - FilteredOut, - HydrationDropKey, - HydrationDropped, - MissingSearchRow, - SearchTraceCollector, - build_fts_page_stage, - build_fusion_stage, - build_rerank_stage, - build_vector_stage, - classify_hydration_drops, - read_manifest_readiness, -) +from basic_memory.repository.search_trace import SearchTraceCollector from basic_memory.repository.semantic_chunking import ( SemanticSourceRow, VectorChunkRecord, @@ -68,7 +59,6 @@ SemanticVectorIndexReconciler, VectorDeletion, VectorKey, - VectorMatch, VectorRecord, ) from basic_memory.repository.semantic_vector_sync import ( @@ -88,45 +78,10 @@ from basic_memory.temporal import TemporalFilter from basic_memory.utils import ensure_timezone_aware -# --- Semantic search constants --- - -VECTOR_FILTER_SCAN_LIMIT = 50000 -# The shared bind-parameter bound for any statement that carries a list of vector -# candidate keys. Both engines cap bind parameters (asyncpg at 32767), so every such -# list — manifest hydration and the filter intersection alike — is split at this size. -VECTOR_HYDRATION_BATCH_SIZE = 250 - -# The manifest conditions under which semantic retrieval will use a stored vector. -# Vector hydration (_hydrate_vector_matches) admits exactly these rows, so anything -# failing them is invisible to search: a chunk left behind by an embedding-model or -# vector-index change, or one still pending. Readiness reporting must apply the same -# predicate — calling such a row "embedded" would report an index settled that -# retrieval cannot answer from, which is the class of lie #1414 exists to remove. -# Callers bind :vector_index and :embedding_model; the scope binds its own IDs. - - -def current_vector_manifest_predicate(scope: ProjectScope, params: dict[str, Any]) -> str: - """SQL admitting only manifest rows retrieval can answer from, within ``scope``.""" - return ( - f"{scope.predicate('project_id', params)} " - "AND vector_index = :vector_index " - "AND embedding_model = :embedding_model " - "AND embedding_status = 'ready'" - ) - +# --- Semantic sync constants --- -# Over-fetch factor for the rerank candidate chunk pool: chunks collapse to unique -# (type, id) rows before reranking, so fetch several times reranker_candidates chunks -# to keep enough unique documents in the rerank window. -RERANK_POOL_CHUNK_FANOUT = 4 -FUSION_BONUS = 0.3 -FUSION_FORMULA_VERSION = "max+0.3*min/v1" -FTS_GATE_THRESHOLD = 0.0 -TOP_CHUNKS_PER_RESULT = 5 -SMALL_NOTE_CONTENT_LIMIT = 2000 OVERSIZED_ENTITY_VECTOR_SHARD_SIZE = semantic_vector_sync.OVERSIZED_ENTITY_VECTOR_SHARD_SIZE _SQLITE_MAX_PREPARE_WINDOW = semantic_vector_sync.SQLITE_MAX_PREPARE_WINDOW -_BUILT_IN_VECTOR_INDEX_NAMES = frozenset({"pgvector", "sqlite-vec"}) type StoredEmbeddingStatus = Literal["pending", "ready"] @@ -216,8 +171,9 @@ class SearchRepositoryBase(ABC): This class defines the common interface that all search repositories must implement, regardless of whether they use SQLite FTS5 or Postgres tsvector for full-text search. - Shared semantic search logic (chunking, embedding orchestration, hybrid score-based fusion) - lives here. Backend-specific operations are delegated to abstract hooks. + Indexing, vector-manifest writes, and embedding orchestration live here. Reading + is delegated to ``SearchReader``, built per call from this repository's current + state so the reader sees the same semantic capability the repository has. Concrete implementations: - SQLiteSearchRepository: Uses FTS5 virtual tables with MATCH queries @@ -335,39 +291,10 @@ async def search( ) -> List[SearchIndexRow]: """Search this repository's project. - ``candidate_keys`` restricts results to those ``(type, id)`` search rows. - ``None`` searches the whole scope; an empty sequence matches nothing. Honored by - the full-text pass, which is where vector and hybrid retrieval evaluate their - structured filters: that pass asks which of a known candidate set a filter - admits instead of paging the filter's whole match set (#1431). - - ``allow_relaxed=True`` retries a zero-result strict multi-word query with - OR-joined content terms. Only the hybrid path opts in: its FTS branch otherwise - contributes nothing for question-form queries. + The reader owns retrieval. This method owns what only the repository knows: + whether its semantic stack is enabled and its vector tables exist. See + ``SearchReader.search`` for ``candidate_keys`` and ``allow_relaxed``. """ - # --- Vector and hybrid: shared retrieval over this repository's scope --- - dispatched = await self._dispatch_retrieval_mode( - search_text=search_text, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - retrieval_mode=retrieval_mode, - min_similarity=min_similarity, - limit=limit, - offset=offset, - trace=trace, - ) - if dispatched is not None: - return dispatched - - # --- Full text: the engine runs the compiled statement --- query = PreparedSearchQuery( search_text=search_text, permalink=permalink, @@ -383,8 +310,8 @@ async def search( retrieval_mode=retrieval_mode, min_similarity=min_similarity, ) - return await self._fts.search( - self.scope, + reader = await self._reader_for(query) + return await reader.search( query, limit=limit, offset=offset, @@ -412,8 +339,6 @@ async def count( allow_relaxed: bool = False, ) -> int: """Count full-text matches with the same filters as ``search``.""" - if retrieval_mode != SearchRetrievalMode.FTS: - raise ValueError("Exact counts are only supported for full-text search retrieval.") query = PreparedSearchQuery( search_text=search_text, permalink=permalink, @@ -429,7 +354,50 @@ async def count( retrieval_mode=retrieval_mode, min_similarity=min_similarity, ) - return await self._fts.count(self.scope, query, allow_relaxed=allow_relaxed) + return await SearchReader(self.scope, self._fts).count(query, allow_relaxed=allow_relaxed) + + # --- Reader construction --- + + def _reranking(self) -> Reranking | None: + """The configured cross-encoder, or None when this repository does not rerank.""" + if self._rerank_provider is None: + return None + return Reranking( + provider=self._rerank_provider, + candidates=self._reranker_candidates, + max_document_chars=self._reranker_max_document_chars, + ) + + def _semantic_search(self) -> SemanticSearch: + """Vector and hybrid retrieval over this repository's live semantic stack. + + Valid once ``_ensure_vector_tables`` has bound the adapter. Read from the + current attributes on every call because the semantic flag can flip at + runtime (#711) and tests retune thresholds between searches. + """ + assert self._embedding_provider is not None + vector = VectorRetrieval( + index=self._semantic_vector_index, + index_name=self._semantic_vector_index_name, + embedding_provider=self._embedding_provider, + embedding_model=self._embedding_model_key(), + vector_k=self._semantic_vector_k, + min_similarity=self._semantic_min_similarity, + ) + return SemanticSearch(self.session_maker, self.scope, self._fts, vector, self._reranking()) + + async def _reader_for(self, query: PreparedSearchQuery) -> SearchReader: + """Build the reader for one call from this repository's current state.""" + # Trigger: the query asks for vector or hybrid retrieval and has text to embed. + # Why: whether semantic search is enabled and whether the vector tables and + # adapter exist is repository lifecycle; the reader only reads. + # Outcome: the semantic gate raises its typed error before any retrieval + # runs, and the reader receives the bound adapter. + if query.retrieval_mode != SearchRetrievalMode.FTS and vector_eligible(query): + self._assert_semantic_available() + await self._ensure_vector_tables() + return SearchReader(self.scope, self._fts, self._semantic_search()) + return SearchReader(self.scope, self._fts) # ------------------------------------------------------------------ # Abstract methods — semantic search (backend-specific DB operations) @@ -440,99 +408,6 @@ async def _ensure_vector_tables(self) -> None: """Create backend-specific vector chunk and embedding tables.""" pass - @logfire.instrument("search.vector_query", extract_args=False) - async def _run_vector_query( - self, - session: AsyncSession, - query_embedding: list[float], - candidate_limit: int, - *, - trace: SearchTraceCollector | None = None, - ) -> list[dict[str, Any]]: - """Query the configured adapter and hydrate only live, ready manifest rows.""" - if trace is not None: - trace.vector = build_vector_stage( - candidate_limit=candidate_limit, - adapter_match_count=0, - hydrated_count=0, - ) - if candidate_limit <= 0: - return [] - - external_vector_index = self._semantic_vector_index_name not in _BUILT_IN_VECTOR_INDEX_NAMES - if not external_vector_index: - matches = await self._semantic_vector_index.search( - query_embedding, - limit=candidate_limit, - ) - if trace is not None: - trace.readiness = await read_manifest_readiness( - session, - self.scope, - self._semantic_vector_index_name, - self._embedding_model_key(), - ) - return await self._hydrate_vector_matches(session, matches, trace=trace) - - scan_limit = min(candidate_limit, VECTOR_FILTER_SCAN_LIMIT) - while True: - matches = await self._semantic_vector_index.search( - query_embedding, - limit=scan_limit, - ) - if trace is not None and trace.readiness is None: - trace.readiness = await read_manifest_readiness( - session, - self.scope, - self._semantic_vector_index_name, - self._embedding_model_key(), - ) - hydrated = await self._hydrate_vector_matches(session, matches, trace=trace) - if ( - len(hydrated) >= candidate_limit - or len(matches) < scan_limit - or scan_limit >= VECTOR_FILTER_SCAN_LIMIT - ): - returned = hydrated[:candidate_limit] - # Trigger: the expanded stale-hit rescan hydrated more chunks than the - # candidate window the search consumes. - # Why: chunks beyond the window never enter thresholding, fusion, or - # reranking — tracing them would invent candidates this execution - # never considered. - # Outcome: the traced stage is trimmed to the returned window. - if trace is not None and trace.vector is not None and len(hydrated) > len(returned): - # Two owners can share one parseable chunk_key (manifest uniqueness - # includes entity_id), so window membership matches by owner too. - returned_chunk_keys = { - (int(row["entity_id"]), str(row["chunk_key"])) for row in returned - } - trimmed: dict[SearchIndexKey, list[tuple[str, float, int | None]]] = {} - for chunk_match in trace.vector.chunk_matches: - if (chunk_match.entity_id, chunk_match.chunk_key) in returned_chunk_keys: - trimmed.setdefault(chunk_match.key, []).append( - ( - chunk_match.chunk_key, - chunk_match.similarity, - chunk_match.entity_id, - ) - ) - # hydrated_count keeps full-scan scope so the vector stage's - # dropped count matches its hydration-drop list; the flattener - # reports the window truncation as its own candidate_window stage. - trace.vector = build_vector_stage( - previous=trace.vector, - chunk_matches=trimmed, - ) - return returned - - # Trigger: stale, pending, or wrong-model adapter hits consumed the - # requested top-k before manifest hydration. - # Why: returning early lets stale extension data crowd every live - # result out of an otherwise valid semantic search. - # Outcome: retry from the same ranked prefix with bounded geometric - # overfetch until enough live rows survive or the adapter is exhausted. - scan_limit = min(scan_limit * 2, VECTOR_FILTER_SCAN_LIMIT) - async def record_entity_vector_deferrals( self, *, @@ -595,110 +470,6 @@ async def record_entity_vector_deferrals( ) await session.commit() - @logfire.instrument("search.vector_manifest_hydration", extract_args=False) - async def _hydrate_vector_matches( - self, - session: AsyncSession, - matches: list[VectorMatch], - *, - trace: SearchTraceCollector | None = None, - ) -> list[dict[str, Any]]: - """Resolve adapter matches through the authoritative ready manifest.""" - if not matches: - return [] - - chunks_by_key: dict[VectorKey, str] = {} - for batch_start in range(0, len(matches), VECTOR_HYDRATION_BATCH_SIZE): - batch = matches[batch_start : batch_start + VECTOR_HYDRATION_BATCH_SIZE] - params: dict[str, Any] = { - "vector_index": self._semantic_vector_index_name, - "embedding_model": self._embedding_model_key(), - } - manifest_predicate = current_vector_manifest_predicate(self.scope, params) - predicates: list[str] = [] - for index, match in enumerate(batch): - params[f"entity_id_{index}"] = match.key.entity_id - params[f"chunk_key_{index}"] = match.key.chunk_key - predicates.append( - f"(entity_id = :entity_id_{index} AND chunk_key = :chunk_key_{index})" - ) - - # Constraint: adapters may return thousands of candidates for deep pages. - # PostgreSQL and SQLite both cap bind parameters, so hydrate in fixed-size - # batches while retaining the adapter's original ranking in the final list. - result = await session.execute( - text( - "SELECT entity_id, chunk_key, chunk_text FROM search_vector_chunks " - "WHERE " + manifest_predicate + " " - "AND (" + " OR ".join(predicates) + ")" - ), - params, - ) - chunks_by_key.update( - { - VectorKey( - entity_id=int(row["entity_id"]), - chunk_key=str(row["chunk_key"]), - ): str(row["chunk_text"]) - for row in result.mappings().all() - } - ) - hydrated = [ - { - "entity_id": match.key.entity_id, - "chunk_key": match.key.chunk_key, - "chunk_text": chunks_by_key[match.key], - "best_similarity": match.similarity, - } - for match in matches - if match.key in chunks_by_key - ] - if trace is not None: - dropped_keys = [ - HydrationDropKey( - entity_id=match.key.entity_id, - chunk_key=match.key.chunk_key, - similarity=match.similarity, - configured_index=self._semantic_vector_index_name, - configured_model=self._embedding_model_key(), - ) - for match in matches - if match.key not in chunks_by_key - ] - drops = await classify_hydration_drops(session, self.scope, dropped_keys) - chunk_matches: dict[SearchIndexKey, list[tuple[str, float, int | None]]] = {} - malformed_drops: list[HydrationDropped] = [] - for row in hydrated: - try: - key = self._parse_chunk_key(str(row["chunk_key"])) - except (ValueError, IndexError): - # A hydrated chunk with an unparseable key silently vanishes from - # retrieval; the trace must name it or the stage counts lie. - malformed_drops.append( - HydrationDropped( - entity_id=int(row["entity_id"]), - chunk_key=str(row["chunk_key"]), - similarity=float(row["best_similarity"]), - reason="malformed_key", - stored_model=None, - stored_index=None, - ) - ) - continue - chunk_matches.setdefault(key, []).append( - (str(row["chunk_key"]), float(row["best_similarity"]), int(row["entity_id"])) - ) - trace.vector = build_vector_stage( - previous=trace.vector, - adapter_match_count=len(matches), - # Malformed keys are dropped, not served — counting them as output - # would contradict the malformed_key rejection listed alongside. - hydrated_count=len(hydrated) - len(malformed_drops), - drops=(*drops, *malformed_drops), - chunk_matches=chunk_matches, - ) - return hydrated - async def _write_embeddings( self, session: AsyncSession, @@ -748,7 +519,7 @@ async def _persist_embeddings( connection = await session.connection() dialect_name = connection.dialect.name external_vector_index = ( - self._semantic_vector_index_name not in _BUILT_IN_VECTOR_INDEX_NAMES + self._semantic_vector_index_name not in BUILT_IN_VECTOR_INDEX_NAMES ) lock_external_write = external_vector_index and dialect_name in {"postgresql", "sqlite"} if external_vector_index: @@ -1027,7 +798,7 @@ async def _finalize_prepared_vector_deletions( for deletion in prepared.staged_deletions ] - external_vector_index = self._semantic_vector_index_name not in _BUILT_IN_VECTOR_INDEX_NAMES + external_vector_index = self._semantic_vector_index_name not in BUILT_IN_VECTOR_INDEX_NAMES if external_vector_index: async with db.scoped_session(self.session_maker) as session: connection = await session.connection() @@ -1076,7 +847,7 @@ async def _finalize_prepared_vector_deletions( return await self._semantic_vector_index.delete(deletions) - if self._semantic_vector_index_name in _BUILT_IN_VECTOR_INDEX_NAMES: + if self._semantic_vector_index_name in BUILT_IN_VECTOR_INDEX_NAMES: return async with db.scoped_session(self.session_maker) as session: await session.execute( @@ -1089,16 +860,6 @@ async def _finalize_prepared_vector_deletions( ) await session.commit() - @abstractmethod - def _distance_to_similarity(self, distance: float) -> float: - """Convert a backend-specific vector distance to cosine similarity in [0, 1]. - - Backend-specific implementations: - - SQLite (vec0): L2/Euclidean distance → cosine similarity via 1 - d²/2 - - Postgres (pgvector <=>): Cosine distance → cosine similarity via 1 - d - """ - pass # pragma: no cover - # ------------------------------------------------------------------ # Shared index / delete operations # ------------------------------------------------------------------ @@ -1393,7 +1154,7 @@ async def _delete_external_entity_vectors_locked( ) -> None: """Delete external vectors after the caller has acquired the project lock.""" self._assert_manifest_vector_ownership(recorded_indexes) - external_indexes = recorded_indexes - _BUILT_IN_VECTOR_INDEX_NAMES + external_indexes = recorded_indexes - BUILT_IN_VECTOR_INDEX_NAMES if not external_indexes: return @@ -1491,14 +1252,14 @@ async def _lock_external_vector_write(self, session: AsyncSession) -> None: def _uses_external_vector_index(self) -> bool: """Return whether this repository writes vectors outside the SQL backend.""" - return self._semantic_vector_index_name not in _BUILT_IN_VECTOR_INDEX_NAMES and hasattr( + return self._semantic_vector_index_name not in BUILT_IN_VECTOR_INDEX_NAMES and hasattr( self, "_semantic_vector_index" ) def _assert_manifest_vector_ownership(self, vector_index_names: Iterable[object]) -> None: """Reject cleanup that cannot reach every externally owned vector.""" recorded_indexes = frozenset(str(name) for name in vector_index_names if str(name)) - external_indexes = recorded_indexes - _BUILT_IN_VECTOR_INDEX_NAMES + external_indexes = recorded_indexes - BUILT_IN_VECTOR_INDEX_NAMES configured_index = self._semantic_vector_index_name if external_indexes and ( not hasattr(self, "_semantic_vector_index") @@ -1565,7 +1326,7 @@ async def _delete_project_vector_rows_in_session( manifest_has_embedding_status = "embedding_status" in manifest_columns configured_index = self._semantic_vector_index_name external_adapter_available = ( - configured_index not in _BUILT_IN_VECTOR_INDEX_NAMES + configured_index not in BUILT_IN_VECTOR_INDEX_NAMES and hasattr(self, "_semantic_vector_index") ) if external_adapter_available: @@ -1606,7 +1367,7 @@ async def _delete_project_vector_rows_in_session( # Outcome: fail before touching any adapter or manifest so the owner can be restored. self._assert_manifest_vector_ownership(entity_ids_by_vector_index) - builtin_indexes = frozenset(entity_ids_by_vector_index) & _BUILT_IN_VECTOR_INDEX_NAMES + builtin_indexes = frozenset(entity_ids_by_vector_index) & BUILT_IN_VECTOR_INDEX_NAMES if manifest_has_embedding_status and (not external_adapter_available or builtin_indexes): builtin_filter = "" if external_adapter_available: @@ -2055,948 +1816,3 @@ def _timestamp_now_expr(self) -> str: SQLite uses CURRENT_TIMESTAMP, Postgres uses NOW(). """ return "CURRENT_TIMESTAMP" - - # ------------------------------------------------------------------ - # Shared semantic search: retrieval mode dispatch - # ------------------------------------------------------------------ - - def _check_vector_eligible( - self, - search_text: Optional[str], - permalink: Optional[str], - permalink_match: Optional[str], - title: Optional[str], - ) -> bool: - """Check whether search_text allows vector / hybrid retrieval.""" - return ( - bool(search_text) - and bool(search_text.strip()) - and search_text.strip() != "*" - and not permalink - and not permalink_match - and not title - ) - - async def _dispatch_retrieval_mode( - self, - *, - search_text: Optional[str], - permalink: Optional[str], - permalink_match: Optional[str], - title: Optional[str], - note_types: Optional[List[str]], - after_date: Optional[datetime], - search_item_types: Optional[List[SearchItemType]], - categories: Optional[List[str]], - metadata_filters: Optional[dict[str, Any]], - file_path_prefix: Optional[str], - temporal: Optional[TemporalFilter], - retrieval_mode: SearchRetrievalMode, - min_similarity: Optional[float] = None, - limit: int, - offset: int, - trace: SearchTraceCollector | None = None, - ) -> Optional[List[SearchIndexRow]]: - """Dispatch vector or hybrid retrieval if requested. - - Returns None when the mode is FTS so the caller should continue - with its backend-specific FTS query. - """ - mode = ( - retrieval_mode.value - if isinstance(retrieval_mode, SearchRetrievalMode) - else str(retrieval_mode) - ) - can_use_vector = self._check_vector_eligible(search_text, permalink, permalink_match, title) - search_text_value = search_text or "" - - if mode == SearchRetrievalMode.VECTOR.value: - if not can_use_vector: - raise ValueError( - "Vector retrieval requires a non-empty text query and does not support " - "title/permalink-only searches." - ) - return await self._search_vector_only( - search_text=search_text_value, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - min_similarity=min_similarity, - limit=limit, - offset=offset, - trace=trace, - ) - if mode == SearchRetrievalMode.HYBRID.value: - if not can_use_vector: - raise ValueError( - "Hybrid retrieval requires a non-empty text query and does not support " - "title/permalink-only searches." - ) - return await self._search_hybrid( - search_text=search_text_value, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - min_similarity=min_similarity, - limit=limit, - offset=offset, - trace=trace, - ) - - # FTS mode: return None to let the subclass handle it - return None - - # ------------------------------------------------------------------ - # Shared semantic search: vector-only retrieval - # ------------------------------------------------------------------ - - @staticmethod - def _parse_chunk_key(chunk_key: str) -> SearchIndexKey: - """Parse a chunk_key like 'observation:5:0' into (type, search_index_id).""" - parts = chunk_key.split(":") - return parts[0], int(parts[1]) - - # ------------------------------------------------------------------ - # Shared semantic search: cross-encoder reranking - # ------------------------------------------------------------------ - - def _should_rerank(self, query_text: str) -> bool: - """Return whether a configured reranker should run for this query.""" - return self._rerank_provider is not None and bool(query_text) - - def _rerank_candidate_limit(self) -> int: - """Return the fixed chunk window that owns reranker-prefix membership.""" - return max( - self._semantic_vector_k, - self._reranker_candidates * RERANK_POOL_CHUNK_FANOUT, - ) - - def _candidate_limit(self, limit: int, offset: int, query_text: str) -> int: - """Size the retrieval candidate *chunk* pool for vector/hybrid search. - - ``candidate_limit`` bounds vector chunks, but many chunks of one large note - collapse to a single ``(type, id)`` row before reranking, so a chunk count does - not equal a unique-document count. When reranking is active we over-fetch by - ``RERANK_POOL_CHUNK_FANOUT`` so a few multi-chunk notes can't starve the rerank - window below ``reranker_candidates`` unique rows. This is best-effort headroom, - not a hard guarantee — a single note dominating the entire nearest-neighbour set - can still yield fewer unique rows (a pathological corpus shape). - """ - if self._should_rerank(query_text): - # Trigger: the requested window extends beyond the fixed reranked prefix. - # Why: a bounded prefix alone can under-fill large pages and hide the - # semantic pagination probe even when more matches exist. - # Outcome: keep prefix membership fixed while adding chunk headroom only - # for the untouched tail that this request must return. - rerank_candidate_limit = self._rerank_candidate_limit() - tail_size = max(0, limit + offset - self._reranker_candidates) - return rerank_candidate_limit + tail_size * 10 - return max(self._semantic_vector_k, (limit + offset) * 10) - - def _rerank_document_text(self, row: SearchIndexRow) -> str: - """Build the document text handed to the cross-encoder for one candidate. - - Prefer the matched chunk (the most relevant passage of a large note), - falling back to the stored snippet. - """ - body = row.matched_chunk_text or row.content_snippet or "" - return build_rerank_document(row.title, body, self._reranker_max_document_chars) - - @staticmethod - def _demote_tail(tail: list[SearchIndexRow], floor: float) -> list[SearchIndexRow]: - """Rescore un-reranked tail rows at or below the floor, preserving their order. - - The reranked pool carries [0, 1] relevance scores while the tail still holds - raw retrieval scores on a different scale ([0, 1.3] for fused hybrid). Left as - is, a tail row could outrank a reranked row numerically. Positive floors put - the tail strictly below the pool; a zero floor yields zeroes because no smaller - score exists in the public [0, 1] range. The returned pool-plus-tail sequence, - rather than a later score-only sort, owns that tie-breaking invariant. - """ - return [ - replace(row, score=score) - for row, score in zip(tail, demote_tail_scores(floor, len(tail))) - ] - - async def _rerank_and_paginate( - self, - query_text: str, - rows: list[SearchIndexRow], - *, - offset: int, - limit: int, - stable_rows: list[SearchIndexRow] | None = None, - trace: SearchTraceCollector | None = None, - ) -> list[SearchIndexRow]: - """Rerank the top candidates, then return the requested ``[offset:offset+limit]`` page. - - Trigger: a reranker is configured and there is a real query. - Why: bi-encoder/FTS ranking lands the gold document in the top-N but often - just below the top-k cutoff (#950); a cross-encoder that reads query and - document together recovers those near-misses. - Outcome: the first ``reranker_candidates`` rows are reordered by reranker - relevance (which replaces ``score``); the requested page is sliced from the - reordered list. - - Every non-empty page rescores the same fixed prefix before slicing so the - untouched tail can be demoted onto the reranker's public ``[0, 1]`` scale. - """ - page_end = offset + limit - if self._rerank_provider is None or not query_text: - return rows[offset:page_end] - - # Trigger: pagination needs more rows than the fixed rerank retrieval window. - # Why: an expanded retrieval may introduce or strengthen raw candidates, but - # letting them replace the original prefix causes duplicates and skips. - # Outcome: the fixed window owns prefix membership; the expanded result only - # supplies new, de-duplicated tail rows. - pool_source = stable_rows if stable_rows is not None else rows - pool = pool_source[: self._reranker_candidates] - pool_keys = {(row.type, row.id) for row in pool} - tail = [row for row in rows if (row.type, row.id) not in pool_keys] - ordered_rows = pool + tail - - # Skip only when there is no prefix to calibrate or the requested page is - # empty. Even a singleton prefix or a wholly-tail page needs the prefix's - # relevance floor so raw hybrid scores cannot leak into cross-project sorting. - if not pool or offset >= len(ordered_rows): - return ordered_rows[offset:page_end] - - pre_rerank_scores = None - if trace is not None: - pre_rerank_scores = {(row.type, row.id): row.score or 0.0 for row in ordered_rows} - documents = [self._rerank_document_text(row) for row in pool] - # A transient provider failure must surface instead of switching this page - # back to retrieval order. A prior page may already have returned reranked - # order, so degrading here can duplicate one result and omit another. - rerank_start = time.perf_counter() if trace is not None else None - with logfire.span( - "search.rerank", - candidate_count=len(pool), - document_chars=sum(map(len, documents)), - ): - scores = validate_rerank_scores( - await self._rerank_provider.rerank(query_text, documents), - len(pool), - ) - - order = sorted(range(len(pool)), key=lambda i: scores[i], reverse=True) - reranked = [replace(pool[i], score=scores[i]) for i in order] - logger.debug( - "Reranked candidates: pool={pool} model={model}", - pool=len(pool), - model=self._rerank_provider.model_name, - ) - tail_floor = reranked[-1].score or 0.0 - demoted_tail = self._demote_tail(tail, floor=tail_floor) - reranked_rows = reranked + demoted_tail - if trace is not None: - assert pre_rerank_scores is not None and rerank_start is not None - trace.rerank = build_rerank_stage( - provider_model=self._rerank_provider.model_name, - reranker_candidates=self._reranker_candidates, - pre_rerank_scores=pre_rerank_scores, - pool_keys=[(row.type, row.id) for row in pool], - rerank_scores={ - (pool[index].type, pool[index].id): score for index, score in enumerate(scores) - }, - post_rerank_rows=[((row.type, row.id), row.score or 0.0) for row in reranked_rows], - demoted_scores={(row.type, row.id): row.score or 0.0 for row in demoted_tail}, - tail_floor=tail_floor, - stable_pool_refetched=trace.stable_pool_refetched, - rerank_ms=(time.perf_counter() - rerank_start) * 1000, - ) - return reranked_rows[offset:page_end] - - async def _search_vector_only( - self, - *, - search_text: str, - permalink: Optional[str], - permalink_match: Optional[str], - title: Optional[str], - note_types: Optional[List[str]], - after_date: Optional[datetime], - search_item_types: Optional[List[SearchItemType]], - categories: Optional[List[str]], - metadata_filters: Optional[dict[str, Any]], - file_path_prefix: Optional[str], - temporal: Optional[TemporalFilter], - min_similarity: Optional[float] = None, - limit: int, - offset: int, - candidate_limit: int | None = None, - _emit_observability_log: bool = True, - _apply_rerank: bool = True, - trace: SearchTraceCollector | None = None, - ) -> List[SearchIndexRow]: - """Run vector-only search returning chunk-level results. - - Returns individual search_index rows (entities, observations, relations) - ranked by vector similarity. Each observation or relation is a first-class - result, not collapsed into its parent entity. - - ``candidate_limit`` is supplied only by a composed retrieval stage that - already sized the shared candidate pool. - """ - self._assert_semantic_available() - await self._ensure_vector_tables() - assert self._embedding_provider is not None - query_text = search_text.strip() - if candidate_limit is None: - candidate_limit = self._candidate_limit(limit, offset, query_text) - query_start = time.perf_counter() - embed_start = time.perf_counter() - with logfire.span("search.embed_query", query_chars=len(query_text)): - query_embedding = await self._embedding_provider.embed_query(query_text) - embed_ms = (time.perf_counter() - embed_start) * 1000 - vector_query_start = time.perf_counter() - - if hasattr(self, "_semantic_vector_index"): - # Constraint: vector adapters may open their own session, while the SQLite - # test/runtime pool can contain only one connection. A plain AsyncSession - # defers checkout until hydration runs after adapter search has released it. - async with self.session_maker() as session: - if trace is None: - vector_rows = await self._run_vector_query( - session, - query_embedding, - candidate_limit, - ) - else: - vector_rows = await self._run_vector_query( - session, - query_embedding, - candidate_limit, - trace=trace, - ) - else: - # Compatibility for focused test repositories that implement the - # pre-extension private query hook without configuring an adapter. - async with db.scoped_session(self.session_maker) as session: - await self._prepare_vector_session(session) - if trace is None: - vector_rows = await self._run_vector_query( - session, - query_embedding, - candidate_limit, - ) - else: - vector_rows = await self._run_vector_query( - session, - query_embedding, - candidate_limit, - trace=trace, - ) - vector_query_ms = (time.perf_counter() - vector_query_start) * 1000 - vector_row_count = len(vector_rows) - hydrate_ms = 0.0 - - if trace is not None: - trace.vector = build_vector_stage( - previous=trace.vector, - effective_min_similarity=( - min_similarity if min_similarity is not None else self._semantic_min_similarity - ), - min_similarity_source=("query" if min_similarity is not None else "config"), - embed_ms=embed_ms, - vector_query_ms=vector_query_ms, - ) - - def _log_vector_summary() -> None: - if not _emit_observability_log: - return - - total_ms = (time.perf_counter() - query_start) * 1000 - if total_ms > 2000: - logger.warning( - "[SEMANTIC_SLOW_QUERY] Semantic query timing: project_id={project_id} " - "retrieval_mode={retrieval_mode} query_length={query_length} " - "candidate_limit={candidate_limit} vector_row_count={vector_row_count} " - "embed_ms={embed_ms:.2f} vector_query_ms={vector_query_ms:.2f} " - "hydrate_ms={hydrate_ms:.2f} total_ms={total_ms:.2f}", - project_id=self.project_id, - retrieval_mode="vector", - query_length=len(query_text), - candidate_limit=candidate_limit, - vector_row_count=vector_row_count, - embed_ms=embed_ms, - vector_query_ms=vector_query_ms, - hydrate_ms=hydrate_ms, - total_ms=total_ms, - ) - - if not vector_rows: - _log_vector_summary() - return [] - - hydrate_start = time.perf_counter() - # Build per-search_index_row similarity scores from chunk-level results. - # Each chunk_key encodes the search_index row type and id; keep both as the - # key because different row types can share the same numeric id (#982). - # Track the best similarity per row (for ranking) and all chunks (for context). - similarity_by_si_key: dict[SearchIndexKey, float] = {} - chunks_by_si_key: dict[SearchIndexKey, list[tuple[float, str]]] = {} - for row in vector_rows: - chunk_key = row.get("chunk_key", "") - if "best_similarity" in row: - similarity = float(row["best_similarity"]) - else: - # Compatibility: private test doubles may still return native distance. - distance = float(row["best_distance"]) - similarity = self._distance_to_similarity(distance) - chunk_text = row.get("chunk_text", "") - try: - si_key = self._parse_chunk_key(chunk_key) - except (ValueError, IndexError): - # Fallback: group by entity_id for chunks without parseable keys - continue - current = similarity_by_si_key.get(si_key) - if current is None or similarity > current: - similarity_by_si_key[si_key] = similarity - chunks_by_si_key.setdefault(si_key, []).append((similarity, chunk_text)) - - if not similarity_by_si_key: - hydrate_ms = (time.perf_counter() - hydrate_start) * 1000 - _log_vector_summary() - return [] - - # Filter out results below the minimum similarity threshold. - # Per-query min_similarity overrides the instance-level default. - effective_min_similarity = ( - min_similarity if min_similarity is not None else self._semantic_min_similarity - ) - if effective_min_similarity > 0.0: - if trace is not None: - threshold_rejections = tuple( - BelowThreshold(key=key, similarity=value, threshold=effective_min_similarity) - for key, value in similarity_by_si_key.items() - if value < effective_min_similarity - ) - trace.vector = build_vector_stage( - previous=trace.vector, - threshold_rejections=threshold_rejections, - ) - similarity_by_si_key = { - k: v for k, v in similarity_by_si_key.items() if v >= effective_min_similarity - } - if not similarity_by_si_key: - hydrate_ms = (time.perf_counter() - hydrate_start) * 1000 - _log_vector_summary() - return [] - - # Fetch the actual search_index rows. Colliding (type, id) keys share one - # bare id, so deduplicate while preserving first-seen order. - si_ids = list(dict.fromkeys(si_id for _, si_id in similarity_by_si_key)) - search_index_rows = await self._fetch_search_index_rows_by_ids(si_ids) - if trace is not None: - trace.vector = build_vector_stage( - previous=trace.vector, - missing_search_rows=tuple( - MissingSearchRow(key=key) - for key in similarity_by_si_key - if key not in search_index_rows - ), - ) - - # Apply optional filters if requested - filter_requested = any( - [ - permalink, - permalink_match, - title, - note_types, - after_date, - search_item_types, - categories, - metadata_filters, - file_path_prefix, - temporal, - ] - ) - - if filter_requested: - allowed_keys = await self._filter_candidate_keys( - list(search_index_rows), - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - ) - if trace is not None: - trace.vector = build_vector_stage( - previous=trace.vector, - filter_rejections=tuple( - FilteredOut(key=key) for key in search_index_rows if key not in allowed_keys - ), - ) - search_index_rows = {k: v for k, v in search_index_rows.items() if k in allowed_keys} - - ranked_rows: list[SearchIndexRow] = [] - for si_key, similarity in similarity_by_si_key.items(): - row = search_index_rows.get(si_key) - if row is None: - continue - - # Small notes: return full content so the answer is always present. - # Large notes: return top-N most relevant chunks for richer context. - content_snippet = row.content_snippet or "" - if content_snippet and len(content_snippet) <= SMALL_NOTE_CONTENT_LIMIT: - matched_chunk_text = content_snippet - else: - si_chunks = chunks_by_si_key.get(si_key, []) - si_chunks.sort(key=lambda c: c[0], reverse=True) - top_texts = [text for _, text in si_chunks[:TOP_CHUNKS_PER_RESULT]] - matched_chunk_text = "\n---\n".join(top_texts) if top_texts else None - - ranked_rows.append( - replace( - row, - score=similarity, - matched_chunk_text=matched_chunk_text, - ) - ) - - ranked_rows.sort(key=lambda item: item.score or 0.0, reverse=True) - hydrate_ms = (time.perf_counter() - hydrate_start) * 1000 - # Rerank over the wide candidate pool, then slice to the page. Suppressed when - # hybrid calls this internally (_apply_rerank=False) — hybrid reranks its own - # fused result; _rerank_and_paginate no-ops back to a plain slice otherwise. - if _apply_rerank: - stable_rows = ranked_rows - if self._should_rerank(query_text): - stable_candidate_limit = self._rerank_candidate_limit() - if candidate_limit > stable_candidate_limit: - if trace is not None: - trace.stable_pool_refetched = True - stable_rows = await self._search_vector_only( - search_text=search_text, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - min_similarity=min_similarity, - limit=stable_candidate_limit, - offset=0, - candidate_limit=stable_candidate_limit, - _emit_observability_log=False, - _apply_rerank=False, - trace=None, - ) - output = await self._rerank_and_paginate( - query_text, - ranked_rows, - offset=offset, - limit=limit, - stable_rows=stable_rows, - trace=trace, - ) - else: - output = ranked_rows[offset : offset + limit] - # Vector latency owns the optional rerank stage too. Logging before the - # awaited provider call hides the feature's dominant cost and can suppress - # the slow-query warning entirely. - _log_vector_summary() - return output - - @logfire.instrument("search.filter_candidates", extract_args=False) - async def _filter_candidate_keys( - self, - candidate_keys: Sequence[SearchIndexKey], - *, - permalink: Optional[str], - permalink_match: Optional[str], - title: Optional[str], - note_types: Optional[List[str]], - after_date: Optional[datetime], - search_item_types: Optional[List[SearchItemType]], - categories: Optional[List[str]], - metadata_filters: Optional[dict[str, Any]], - file_path_prefix: Optional[str], - temporal: Optional[TemporalFilter], - ) -> set[SearchIndexKey]: - """Return which of ``candidate_keys`` the structured filters admit. - - Vector retrieval scores embeddings and cannot evaluate a structured filter, so - the surviving candidates are decided by an FTS-mode pass carrying every filter. - Asking that pass for a *page of the filter's whole match set* and intersecting - client-side silently lost any candidate that sorted past the page (#1431); asking - it about the candidates themselves cannot, because the answer is bounded by the - question. - - The candidate list is split at the shared bind-parameter bound, so a deep page - whose candidate pool runs to thousands of rows costs a few small indexed lookups - instead of one unbounded scan. - """ - allowed_keys: set[SearchIndexKey] = set() - for batch_start in range(0, len(candidate_keys), VECTOR_HYDRATION_BATCH_SIZE): - batch = candidate_keys[batch_start : batch_start + VECTOR_HYDRATION_BATCH_SIZE] - filtered_rows = await self.search( - search_text=None, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - retrieval_mode=SearchRetrievalMode.FTS, - # The restriction, not this limit, is what bounds the result: one row per - # requested key, since (id, type, project_id) identifies a search row. - limit=len(batch), - offset=0, - candidate_keys=batch, - ) - allowed_keys.update((row.type, row.id) for row in filtered_rows if row.id is not None) - return allowed_keys - - @logfire.instrument("search.fetch_candidate_rows", extract_args=False) - async def _fetch_search_index_rows_by_ids( - self, row_ids: list[int] - ) -> dict[SearchIndexKey, SearchIndexRow]: - """Fetch search_index rows by id, keyed by (type, id) to disambiguate types. - - A bare id can match one row per type (independent id sequences), so the - result must carry every matching row rather than letting one clobber another. - """ - if not row_ids: - return {} - placeholders = ",".join(f":id_{idx}" for idx in range(len(row_ids))) - params: dict[str, Any] = {f"id_{idx}": rid for idx, rid in enumerate(row_ids)} - scope_predicate = self.scope.predicate("project_id", params) - sql = f""" - SELECT - project_id, id, title, permalink, file_path, type, metadata, - from_id, to_id, relation_type, entity_id, content_snippet, - category, created_at, updated_at, 0 as score - FROM search_index - WHERE {scope_predicate} - AND id IN ({placeholders}) - """ - result: dict[SearchIndexKey, SearchIndexRow] = {} - async with db.scoped_session(self.session_maker) as session: - row_result = await session.execute(text(sql), params) - for row in row_result.fetchall(): - search_row = SearchIndexRow.from_mapping(row._asdict()) - result[(search_row.type, search_row.id)] = search_row - return result - - # ------------------------------------------------------------------ - # Shared semantic search: hybrid score-based fusion - # ------------------------------------------------------------------ - - async def _search_hybrid( - self, - *, - search_text: str, - permalink: Optional[str], - permalink_match: Optional[str], - title: Optional[str], - note_types: Optional[List[str]], - after_date: Optional[datetime], - search_item_types: Optional[List[SearchItemType]], - categories: Optional[List[str]], - metadata_filters: Optional[dict[str, Any]], - file_path_prefix: Optional[str], - temporal: Optional[TemporalFilter], - min_similarity: Optional[float] = None, - limit: int, - offset: int, - _candidate_limit_override: int | None = None, - _apply_rerank: bool = True, - _emit_observability_log: bool = True, - trace: SearchTraceCollector | None = None, - ) -> List[SearchIndexRow]: - """Fuse FTS and vector results using score-based fusion. - - Uses the search_index (type, id) pair as the fusion key. The formula - ``max(vec, fts) + FUSION_BONUS * min(vec, fts)`` preserves - the dominant signal and rewards dual-source agreement. - """ - self._assert_semantic_available() - query_text = search_text.strip() - rerank_configured = self._should_rerank(query_text) - rerank_enabled = _apply_rerank and rerank_configured - query_start = time.perf_counter() - candidate_limit = ( - _candidate_limit_override - if _candidate_limit_override is not None - else self._candidate_limit(limit, offset, query_text) - ) - fts_start = time.perf_counter() - # allow_relaxed: question-form queries rarely AND-match, and a dead FTS - # branch silently degrades hybrid to vector-only ranking. Fusion plus - # bm25 keep relaxed lexical candidates from dominating precision. - with logfire.span("search.fts", candidate_limit=candidate_limit) as fts_span: - fts_results = await self.search( - search_text=search_text, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - retrieval_mode=SearchRetrievalMode.FTS, - limit=candidate_limit, - offset=0, - allow_relaxed=True, - trace=trace, - ) - fts_span.set_attribute("result_count", len(fts_results)) - fts_ms = (time.perf_counter() - fts_start) * 1000 - vector_start = time.perf_counter() - vector_results = await self._search_vector_only( - search_text=search_text, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - min_similarity=min_similarity, - limit=candidate_limit, - offset=0, - # Trigger: reranking owns a bounded candidate window shared by both legs. - # Why: the disabled path historically expands the vector leg again to - # preserve recall when many vector chunks collapse into a few search rows. - # Outcome: avoid double expansion only when reranking is actually active. - candidate_limit=candidate_limit if rerank_configured else None, - _emit_observability_log=False, - _apply_rerank=False, - trace=trace, - ) - vector_ms = (time.perf_counter() - vector_start) * 1000 - # Trigger: with reranking disabled the vector leg expands internally and can - # hydrate more rows than the fusion window it returns. - # Why: rows cut here never fuse — left in the trace they would surface as - # candidates with no rejection and no fused rank, which the response labels - # "returned". Rows with a recorded rejection keep their chunk evidence. - # Outcome: the trace keeps rows handed to fusion (or explicitly rejected); - # the cut shows up as served-chunk shrinkage in the candidate_window stage. - if trace is not None and trace.vector is not None: - kept_row_keys = {(row.type, row.id) for row in vector_results} - kept_row_keys.update( - rejection.key - for rejection_group in ( - trace.vector.threshold_rejections, - trace.vector.filter_rejections, - trace.vector.missing_search_rows, - ) - for rejection in rejection_group - ) - if any(match.key not in kept_row_keys for match in trace.vector.chunk_matches): - fused_chunks: dict[SearchIndexKey, list[tuple[str, float, int | None]]] = {} - for chunk_match in trace.vector.chunk_matches: - if chunk_match.key in kept_row_keys: - fused_chunks.setdefault(chunk_match.key, []).append( - (chunk_match.chunk_key, chunk_match.similarity, chunk_match.entity_id) - ) - trace.vector = build_vector_stage( - previous=trace.vector, - chunk_matches=fused_chunks, - ) - fusion_start = time.perf_counter() - - with logfire.span( - "search.fusion", fts_count=len(fts_results), vector_count=len(vector_results) - ) as fusion_span: - # --- Score-based fusion keyed on (type, id) --- - # A bare row id collides across row types (independent id sequences), so - # fusion must key on (type, id) or distinct rows would merge (#982). - # FTS scores are normalized to [0, 1] (BM25 is unbounded). - # Vector scores are used raw — already calibrated [0, 1] by _distance_to_similarity(). - rows_by_key: dict[SearchIndexKey, SearchIndexRow] = {} - - # Normalize FTS scores to [0, 1] — handles both SQLite (negative bm25) - # and Postgres (positive ts_rank) by using absolute values - fts_abs = [abs(row.score or 0.0) for row in fts_results] - fts_max = max(fts_abs) if fts_abs else 1.0 - - fts_scores: dict[SearchIndexKey, float] = {} - fts_ranks: dict[SearchIndexKey, int] = {} - for rank, row in enumerate(fts_results): - if row.id is None: - continue - row_key = (row.type, row.id) - norm = abs(row.score or 0.0) / fts_max if fts_max > 0 else 0.0 - # Gate: FTS scores below threshold contribute zero - if norm < FTS_GATE_THRESHOLD: - norm = 0.0 - fts_scores[row_key] = norm - fts_ranks.setdefault(row_key, rank) - rows_by_key[row_key] = row - - if trace is not None: - relaxed_fallback_used = ( - trace.fts.relaxed_fallback_used if trace.fts is not None else False - ) - trace.fts = build_fts_page_stage( - [((row.type, row.id), row.score or 0.0) for row in fts_results], - normalized_scores=fts_scores, - entity_ids={(row.type, row.id): row.entity_id for row in fts_results}, - fts_max_abs=fts_max, - relaxed_fallback_used=relaxed_fallback_used, - fts_ms=fts_ms, - ) - - vec_scores: dict[SearchIndexKey, float] = {} - vec_ranks: dict[SearchIndexKey, int] = {} - for rank, row in enumerate(vector_results): - if row.id is None: - continue - row_key = (row.type, row.id) - # Trigger: no re-normalization by vec_max - # Why: vector similarity is already calibrated [0, 1]; re-normalizing - # inflates weak matches when the entire result set is mediocre - vec_scores[row_key] = row.score or 0.0 - vec_ranks.setdefault(row_key, rank) - rows_by_key[row_key] = row - - # Fuse: max(v, f) + FUSION_BONUS * min(v, f) - # Preserves the dominant signal; bonus rewards dual-source agreement. - # Output range: [0, 1.3] for dual-source, [0, 1.0] for single-source. - fused_scores: dict[SearchIndexKey, float] = {} - for row_key in fts_scores.keys() | vec_scores.keys(): - v = vec_scores.get(row_key, 0.0) - f = fts_scores.get(row_key, 0.0) - fused_scores[row_key] = max(v, f) + FUSION_BONUS * min(v, f) - - ranked = sorted(fused_scores.items(), key=lambda item: item[1], reverse=True) - fusion_span.set_attribute("result_count", len(ranked)) - fusion_ms = (time.perf_counter() - fusion_start) * 1000 - if trace is not None: - trace.fusion = build_fusion_stage( - formula_version=FUSION_FORMULA_VERSION, - bonus=FUSION_BONUS, - fts_scores=fts_scores, - fts_ranks=fts_ranks, - vector_scores=vec_scores, - vector_ranks=vec_ranks, - ranked_scores=ranked, - fusion_ms=fusion_ms, - ) - - def _materialize(entry: tuple[SearchIndexKey, float]) -> SearchIndexRow: - row_key, fused_score = entry - row = rows_by_key[row_key] - # FTS-only hits use the bounded content preview and its truncation metadata. - # Copying the full note into matched_chunk bypasses that response bound. - return replace(row, score=fused_score) - - # Rerank the top fused candidates before paginating. When reranking is active - # we materialize the whole candidate list (cheap next to a cross-encoder call) - # and hand it to the shared paginate helper; the disabled path stays cheap by - # materializing only the requested page. - if rerank_enabled: - candidates = [_materialize(entry) for entry in ranked] - stable_candidates = candidates - stable_candidate_limit = self._rerank_candidate_limit() - if candidate_limit > stable_candidate_limit: - if trace is not None: - trace.stable_pool_refetched = True - stable_candidates = await self._search_hybrid( - search_text=search_text, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - categories=categories, - metadata_filters=metadata_filters, - file_path_prefix=file_path_prefix, - temporal=temporal, - min_similarity=min_similarity, - limit=stable_candidate_limit, - offset=0, - _candidate_limit_override=stable_candidate_limit, - _apply_rerank=False, - _emit_observability_log=False, - trace=None, - ) - stable_keys = {(row.type, row.id) for row in stable_candidates} - expanded_tail = [entry for entry in ranked if entry[0] not in stable_keys] - - # Trigger: deeper pages expand the FTS/vector retrieval windows. - # Why: score fusion can strengthen an existing row when its second - # signal appears later, moving it across a page already returned. - # Outcome: freeze the fixed fused universe, then order newly admitted - # rows by their earliest source rank. That rank cannot improve after a - # row first appears, so each larger window only appends to the tail. - expanded_tail.sort( - key=lambda entry: ( - min( - fts_ranks.get(entry[0], candidate_limit), - vec_ranks.get(entry[0], candidate_limit), - ), - entry[0], - ) - ) - candidates = stable_candidates + [_materialize(entry) for entry in expanded_tail] - output = await self._rerank_and_paginate( - query_text, - candidates, - offset=offset, - limit=limit, - stable_rows=stable_candidates, - trace=trace, - ) - else: - output = [_materialize(entry) for entry in ranked[offset : offset + limit]] - total_ms = (time.perf_counter() - query_start) * 1000 - if _emit_observability_log and total_ms > 2500: - logger.warning( - "[SEMANTIC_SLOW_QUERY] Semantic query timing: project_id={project_id} " - "retrieval_mode={retrieval_mode} query_length={query_length} " - "candidate_limit={candidate_limit} fts_count={fts_count} " - "vector_count={vector_count} fts_ms={fts_ms:.2f} vector_ms={vector_ms:.2f} " - "fusion_ms={fusion_ms:.2f} total_ms={total_ms:.2f}", - project_id=self.project_id, - retrieval_mode="hybrid", - query_length=len(query_text), - candidate_limit=candidate_limit, - fts_count=len(fts_results), - vector_count=len(vector_results), - fts_ms=fts_ms, - vector_ms=vector_ms, - fusion_ms=fusion_ms, - total_ms=total_ms, - ) - return output diff --git a/src/basic_memory/repository/sqlite_search_repository.py b/src/basic_memory/repository/sqlite_search_repository.py index aa75461e7..0861023b7 100644 --- a/src/basic_memory/repository/sqlite_search_repository.py +++ b/src/basic_memory/repository/sqlite_search_repository.py @@ -407,15 +407,6 @@ async def drop_vector_tables(self) -> None: await session.commit() self._vector_tables_initialized = False - @override - def _distance_to_similarity(self, distance: float) -> float: - """Convert L2 distance to cosine similarity for normalized embeddings. - - sqlite-vec vec0 returns Euclidean (L2) distance by default. - For unit-normalized vectors: L2² = 2·(1 - cos_sim), so cos_sim = 1 - L2²/2. - """ - return max(0.0, 1.0 - (distance * distance) / 2.0) - @asynccontextmanager @override async def _prepare_entity_write_scope(self): diff --git a/src/basic_memory/services/project_readiness.py b/src/basic_memory/services/project_readiness.py index ed7bc1ca8..006bf0899 100644 --- a/src/basic_memory/services/project_readiness.py +++ b/src/basic_memory/services/project_readiness.py @@ -24,7 +24,7 @@ from basic_memory.repository.embedding_provider_factory import ( configured_embedding_provider_identity, ) -from basic_memory.repository.search_repository_base import current_vector_manifest_predicate +from basic_memory.repository.search_reader import current_vector_manifest_predicate from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_vector_index_factory import ( resolve_semantic_vector_index_name, diff --git a/src/basic_memory/services/retrieval_inspect.py b/src/basic_memory/services/retrieval_inspect.py index b57bb771e..e70c17889 100644 --- a/src/basic_memory/services/retrieval_inspect.py +++ b/src/basic_memory/services/retrieval_inspect.py @@ -16,11 +16,8 @@ from basic_memory.repository.note_content_repository import NoteContentRepository from basic_memory.repository.search_index_row import SearchIndexRow from basic_memory.repository.search_repository import SearchRepository -from basic_memory.repository.search_repository_base import ( - ChunkManifestRow, - FUSION_FORMULA_VERSION, - SearchRepositoryBase, -) +from basic_memory.repository.search_reader import FUSION_FORMULA_VERSION, parse_chunk_key +from basic_memory.repository.search_repository_base import ChunkManifestRow from basic_memory.repository.search_trace import ( FinalResultEntry, QueryMeta, @@ -416,7 +413,7 @@ async def inspect_entity_chunks( inspected_chunks: list[InspectedChunk] = [] chunks_by_search_row: dict[tuple[str, int], list[InspectedChunk]] = {} for stored_row in stored_rows: - row_key = SearchRepositoryBase._parse_chunk_key(stored_row.chunk_key) + row_key = parse_chunk_key(stored_row.chunk_key) ordinal = int(stored_row.chunk_key.split(":")[2]) inspected_chunk = InspectedChunk( stored_row=stored_row, diff --git a/test-int/semantic/test_multilingual_benchmark_contract.py b/test-int/semantic/test_multilingual_benchmark_contract.py index 435562f03..461ddafa2 100644 --- a/test-int/semantic/test_multilingual_benchmark_contract.py +++ b/test-int/semantic/test_multilingual_benchmark_contract.py @@ -18,7 +18,7 @@ create_sqlite_search_vector_embeddings, ) from basic_memory.repository.semantic_chunking import split_text_into_chunks -from basic_memory.repository.search_repository_base import SMALL_NOTE_CONTENT_LIMIT +from basic_memory.repository.search_reader import SMALL_NOTE_CONTENT_LIMIT from basic_memory.schemas.search import SearchRetrievalMode from semantic.multilingual_benchmark import ( diff --git a/test-int/semantic/test_search_diagnostics.py b/test-int/semantic/test_search_diagnostics.py index 3179ed9d1..e7583fad7 100644 --- a/test-int/semantic/test_search_diagnostics.py +++ b/test-int/semantic/test_search_diagnostics.py @@ -340,20 +340,15 @@ async def test_similarity_formula_analysis(sqlite_engine_factory, tmp_path): from basic_memory import db as bm_db repo = cast(Any, service.repository) + semantic = repo._semantic_search() async with bm_db.scoped_session(repo.session_maker) as session: - await repo._prepare_vector_session(session) - vector_rows = await repo._run_vector_query( - session, - query_embedding, - candidate_limit=20, - ) + vector_rows = await semantic._run_vector_query(session, query_embedding, 20) print(f"\nQuery: '{query_text}'") print(f" {'chunk_key':<40} {'similarity':>12}") - for row in vector_rows[:10]: - similarity = float(row["best_similarity"]) - assert 0.0 <= similarity <= 1.0 - print(f" {row['chunk_key']:<40} {similarity:>12.4f}") + for chunk in vector_rows[:10]: + assert 0.0 <= chunk.similarity <= 1.0 + print(f" {chunk.chunk_key:<40} {chunk.similarity:>12.4f}") # --- Test: min_similarity threshold effectiveness --- diff --git a/test-int/semantic/test_semantic_coverage.py b/test-int/semantic/test_semantic_coverage.py index 0e4881089..d71faaca0 100644 --- a/test-int/semantic/test_semantic_coverage.py +++ b/test-int/semantic/test_semantic_coverage.py @@ -3,7 +3,7 @@ Exercises the uncovered code paths in PostgresSearchRepository: - _ensure_vector_tables (lines 258-352): pgvector extension, table creation, dimension mismatch detection -- _run_vector_query (lines 389-429): vector similarity query with cosine distance +- SemanticSearch._run_vector_query: vector similarity query with cosine distance - _write_embeddings (lines 431-458): embedding upsert into pgvector table - Metadata filters in FTS search (lines 682-745): JSONB filter operators (eq, in, contains, gt/gte/lt/lte, between) @@ -19,6 +19,7 @@ from basic_memory import db from basic_memory.config import DatabaseBackend +from basic_memory.repository.search_reader import HydratedChunk, SemanticSearch from basic_memory.schemas.search import SearchItemType, SearchQuery, SearchRetrievalMode from semantic.conftest import ( @@ -104,7 +105,7 @@ async def test_postgres_vector_table_setup_and_query(postgres_engine_factory, tm async def test_postgres_hybrid_search(postgres_engine_factory, tmp_path): """Exercise the hybrid (score-based fusion) code path on Postgres. - This covers the full _search_hybrid path including both FTS and vector + This covers the full SemanticSearch.hybrid path including both FTS and vector retrieval with score-based fusion. """ skip_if_needed(PG_FASTEMBED) @@ -118,7 +119,7 @@ async def test_postgres_hybrid_search(postgres_engine_factory, tmp_path): await seed_benchmark_notes(search_service, note_count=20) - # Hybrid search — exercises _search_hybrid score-based fusion + # Hybrid search — exercises SemanticSearch.hybrid score-based fusion results = await search_service.search( SearchQuery( text="database migration schema", @@ -158,17 +159,20 @@ async def test_postgres_hybrid_preserves_candidate_windows( repo._reranker_candidates = 100 candidate_limits: list[int] = [] - run_vector_query = repo._run_vector_query + run_vector_query = SemanticSearch._run_vector_query async def record_vector_query( + self: SemanticSearch, session: Any, query_embedding: list[float], candidate_limit: int, - ) -> list[dict[str, Any]]: + *, + trace: Any = None, + ) -> list[HydratedChunk]: candidate_limits.append(candidate_limit) - return await run_vector_query(session, query_embedding, candidate_limit) + return await run_vector_query(self, session, query_embedding, candidate_limit, trace=trace) - monkeypatch.setattr(repo, "_run_vector_query", record_vector_query) + monkeypatch.setattr(SemanticSearch, "_run_vector_query", record_vector_query) baseline_results = await search_service.search( SearchQuery( diff --git a/tests/repository/test_distance_to_similarity.py b/tests/repository/test_distance_to_similarity.py deleted file mode 100644 index faa8a27c0..000000000 --- a/tests/repository/test_distance_to_similarity.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Unit tests for backend-specific distance-to-similarity conversions.""" - -import pytest - -from basic_memory.repository.postgres_search_repository import PostgresSearchRepository -from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository - - -def test_sqlite_distance_to_similarity_formula(): - """SQLite converts L2 distance to cosine similarity for normalized vectors.""" - repo = SQLiteSearchRepository.__new__(SQLiteSearchRepository) - - assert repo._distance_to_similarity(0.0) == 1.0 - assert repo._distance_to_similarity(1.0) == pytest.approx(0.5) - assert repo._distance_to_similarity(2.0) == 0.0 - - -def test_postgres_distance_to_similarity_formula(): - """Postgres converts pgvector cosine distance to cosine similarity.""" - repo = PostgresSearchRepository.__new__(PostgresSearchRepository) - - assert repo._distance_to_similarity(0.0) == 1.0 - assert repo._distance_to_similarity(1.0) == 0.0 - assert repo._distance_to_similarity(2.0) == 0.0 diff --git a/tests/repository/test_hybrid_fusion.py b/tests/repository/test_hybrid_fusion.py index 0ad16896b..dd5e14dc4 100644 --- a/tests/repository/test_hybrid_fusion.py +++ b/tests/repository/test_hybrid_fusion.py @@ -6,26 +6,26 @@ 3. Produces zero fused score when the source score is zero """ -from sqlalchemy.ext.asyncio import AsyncSession -from basic_memory.repository.search_scope import ProjectScope -from collections.abc import Sequence +from collections.abc import Callable, Sequence from dataclasses import dataclass -from datetime import datetime -from typing import override, Any, Optional, cast +from typing import Any, cast from unittest.mock import AsyncMock, patch import pytest +from sqlalchemy.ext.asyncio import AsyncSession from basic_memory.repository.embedding_provider import EmbeddingProvider -from basic_memory.repository.search_index_row import SearchIndexRow -from basic_memory.repository.search_repository_base import ( +from basic_memory.repository.search_index_row import SearchIndexKey, SearchIndexRow +from basic_memory.repository.search_query import PreparedSearchQuery +from basic_memory.repository.search_reader import ( FUSION_BONUS, - SearchIndexKey, - SearchRepositoryBase, + SemanticSearch, + VectorRetrieval, ) +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.search_trace import SearchTraceCollector -from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode -from basic_memory.temporal import TemporalFilter +from basic_memory.repository.semantic_vector_index import SemanticVectorIndex +from basic_memory.schemas.search import SearchRetrievalMode @dataclass @@ -51,159 +51,108 @@ class FakeRow: matched_chunk_text: str | None = None -class ConcreteSearchRepo(SearchRepositoryBase): - """Minimal concrete subclass for testing hybrid fusion logic.""" +class FakeFts: + """An ``FtsBackend`` that answers every pass with fixed rows and records what it was asked. - def __init__(self): - self._semantic_enabled = True - self._semantic_vector_k = 100 - self._semantic_min_similarity = 0.0 - # _search_hybrid calls _assert_semantic_available which checks this - self._embedding_provider = _fake_embedding_provider() - self._vector_dimensions = 384 - self._vector_tables_initialized = True - self.session_maker = None - self.project_id = 1 - self.scope = ProjectScope.single(1) - - @override - async def init_search_index(self): - pass # pragma: no cover + ``rows`` is the answer, or a function of the requested ``limit`` when a test needs + the lexical leg to widen with the candidate window. + """ - @override - async def get_entity_physical_chunk_keys(self, entity_id: int) -> set[str] | None: - return None # physical storage is not inspectable in this double + def __init__(self, rows: Sequence[Any] | Callable[[int], Sequence[Any]] = ()) -> None: + if isinstance(rows, Sequence): + fixed = list(rows) + self.answer: Callable[[int], Sequence[Any]] = lambda _limit: fixed + else: + self.answer = rows + self.queries: list[PreparedSearchQuery] = [] + self.calls: list[dict[str, Any]] = [] - @override async def search( self, - search_text: Optional[str] = None, - permalink: Optional[str] = None, - permalink_match: Optional[str] = None, - title: Optional[str] = None, - note_types: Optional[list[str]] = None, - after_date: Optional[datetime] = None, - search_item_types: Optional[list[SearchItemType]] = None, - categories: Optional[list[str]] = None, - metadata_filters: Optional[dict[str, Any]] = None, - file_path_prefix: Optional[str] = None, - temporal: Optional[TemporalFilter] = None, - retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, - min_similarity: Optional[float] = None, - limit: int = 10, - offset: int = 0, + scope: ProjectScope, + query: PreparedSearchQuery, + *, + limit: int, + offset: int, allow_relaxed: bool = False, session: AsyncSession | None = None, - *, candidate_keys: Sequence[SearchIndexKey] | None = None, trace: SearchTraceCollector | None = None, ) -> list[SearchIndexRow]: - return [] # pragma: no cover - - @override - async def _ensure_vector_tables(self): - pass # pragma: no cover - - @override - async def _run_vector_query( - self, - session, - query_embedding, - candidate_limit, - *, - trace: SearchTraceCollector | None = None, - ): - return [] # pragma: no cover - - @override - async def _write_embeddings(self, session, jobs, embeddings): - pass # pragma: no cover - - @override - async def _delete_entity_chunks(self, session, entity_id, *, expected_deletions=None): - return [] # pragma: no cover - - @override - async def _delete_stale_chunks( + self.queries.append(query) + self.calls.append( + { + "limit": limit, + "offset": offset, + "allow_relaxed": allow_relaxed, + "candidate_keys": candidate_keys, + } + ) + return cast(list[SearchIndexRow], list(self.answer(limit))) + + async def count( self, - session, - stale_ids, - entity_id, + scope: ProjectScope, + query: PreparedSearchQuery, *, - expected_deletions=None, - ): - return [] # pragma: no cover + allow_relaxed: bool = False, + ) -> int: + return len(self.answer(0)) + + +def fake_vector_retrieval( + *, + vector_k: int = 100, + min_similarity: float = 0.0, + embed_query: AsyncMock | None = None, +) -> VectorRetrieval: + """A semantic stack whose adapter is never consulted: tests stub the neighbour stage.""" + provider = type( + "EP", + (), + { + "model_name": "fake", + "dimensions": 384, + "embed_query": embed_query or AsyncMock(return_value=[0.0] * 384), + "embed_documents": AsyncMock(return_value=[]), + "runtime_log_attrs": lambda self: {}, + }, + )() + return VectorRetrieval( + index=cast(SemanticVectorIndex, object()), + index_name="sqlite-vec", + embedding_provider=cast(EmbeddingProvider, provider), + embedding_model="fake:384", + vector_k=vector_k, + min_similarity=min_similarity, + ) - async def _update_timestamp_sql(self): - return "CURRENT_TIMESTAMP" # pragma: no cover - @override - def _distance_to_similarity(self, distance: float) -> float: - return 1.0 / (1.0 + max(distance, 0.0)) # pragma: no cover +HYBRID_QUERY = PreparedSearchQuery(search_text="test", retrieval_mode=SearchRetrievalMode.HYBRID) -def _fake_embedding_provider() -> EmbeddingProvider: - return cast( - EmbeddingProvider, - type( - "EP", - (), - { - "model_name": "fake", - "dimensions": 384, - "embed_query": AsyncMock(return_value=[0.0] * 384), - "embed_documents": AsyncMock(return_value=[]), - "runtime_log_attrs": lambda self: {}, - }, - )(), +async def fuse( + fts_results: list[Any], + vector_results: list[Any], + *, + query: PreparedSearchQuery = HYBRID_QUERY, +) -> list[SearchIndexRow]: + """Run hybrid with both legs answering fixed rows, so only fusion is under test.""" + semantic = SemanticSearch( + cast(Any, None), ProjectScope.single(1), FakeFts(fts_results), fake_vector_retrieval() ) - - -HYBRID_KWARGS: dict[str, Any] = dict( - search_text="test", - permalink=None, - permalink_match=None, - title=None, - note_types=None, - after_date=None, - search_item_types=None, - categories=None, - metadata_filters=None, - file_path_prefix=None, - temporal=None, - limit=10, - offset=0, -) + with patch.object(semantic, "vector_only", new_callable=AsyncMock, return_value=vector_results): + return await semantic.hybrid(query, limit=10, offset=0) @pytest.mark.asyncio async def test_high_fts_score_boosts_ranking(): """FTS-only: a high normalized score should outscore a low normalized score.""" - repo = ConcreteSearchRepo() - - # Two FTS results with very different scores high_score_row = FakeRow(id=1, score=10.0, title="high") low_score_row = FakeRow(id=2, score=0.5, title="low") - fts_results = [high_score_row, low_score_row] # No vector results — isolate FTS weighting behavior - vector_results = [] - - with ( - patch.object( - repo, - "search", - new_callable=AsyncMock, - return_value=fts_results, - ), - patch.object( - repo, - "_search_vector_only", - new_callable=AsyncMock, - return_value=vector_results, - ), - ): - results = await repo._search_hybrid(**HYBRID_KWARGS) + results = await fuse([high_score_row, low_score_row], []) assert len(results) == 2 # After normalization: id=1 → 1.0, id=2 → 0.05 @@ -215,8 +164,6 @@ async def test_high_fts_score_boosts_ranking(): @pytest.mark.asyncio async def test_dual_source_ranks_higher_than_single(): """A result in both FTS and vector should rank above single-source results.""" - repo = ConcreteSearchRepo() - # Row 1 in both (fts=5.0→norm 1.0, vec=0.9), Row 2 FTS-only (fts=5.0→norm 1.0), # Row 3 vec-only (0.8) fts_results = [ @@ -228,13 +175,7 @@ async def test_dual_source_ranks_higher_than_single(): FakeRow(id=3, score=0.8, title="vec-only"), ] - with ( - patch.object(repo, "search", new_callable=AsyncMock, return_value=fts_results), - patch.object( - repo, "_search_vector_only", new_callable=AsyncMock, return_value=vector_results - ), - ): - results = await repo._search_hybrid(**HYBRID_KWARGS) + results = await fuse(fts_results, vector_results) result_ids = [r.id for r in results] # Row 1 (dual-source) should rank first, then Row 2 (FTS 1.0), then Row 3 (vec 0.8) @@ -251,19 +192,7 @@ async def test_dual_source_ranks_higher_than_single(): @pytest.mark.asyncio async def test_zero_score_produces_zero_fused(): """A zero-score FTS result with no vector match produces a zero fused score.""" - repo = ConcreteSearchRepo() - - # FTS result with score 0.0 - fts_results = [FakeRow(id=1, score=0.0, title="zero-score")] - vector_results = [] - - with ( - patch.object(repo, "search", new_callable=AsyncMock, return_value=fts_results), - patch.object( - repo, "_search_vector_only", new_callable=AsyncMock, return_value=vector_results - ), - ): - results = await repo._search_hybrid(**HYBRID_KWARGS) + results = await fuse([FakeRow(id=1, score=0.0, title="zero-score")], []) assert len(results) == 1 # Zero FTS score, no vector → fused = max(0, 0) + 0.3 * min(0, 0) = 0.0 @@ -277,18 +206,10 @@ async def test_cross_type_id_collision_keeps_both_results(): search_index row types have independent id sequences, so fusing on a bare row id merged unrelated rows into one result and dropped the other. """ - repo = ConcreteSearchRepo() - fts_results = [FakeRow(id=1, type="entity", score=5.0, title="entity-row")] vector_results = [FakeRow(id=1, type="relation", score=0.8, title="relation-row")] - with ( - patch.object(repo, "search", new_callable=AsyncMock, return_value=fts_results), - patch.object( - repo, "_search_vector_only", new_callable=AsyncMock, return_value=vector_results - ), - ): - results = await repo._search_hybrid(**HYBRID_KWARGS) + results = await fuse(fts_results, vector_results) assert {(r.type, r.id) for r in results} == {("entity", 1), ("relation", 1)} # Single-source scores must not earn the dual-source fusion bonus across types. @@ -301,21 +222,9 @@ async def test_cross_type_id_collision_keeps_both_results(): @pytest.mark.asyncio async def test_fts_only_result_does_not_copy_content_into_matched_chunk(): """FTS-only hits use the API content preview instead of a second full-note field.""" - repo = ConcreteSearchRepo() - content = "This is the full note content with the answer we need to find." - fts_results = [ - FakeRow(id=1, score=5.0, title="fts-hit", content_snippet=content), - ] - vector_results = [] - with ( - patch.object(repo, "search", new_callable=AsyncMock, return_value=fts_results), - patch.object( - repo, "_search_vector_only", new_callable=AsyncMock, return_value=vector_results - ), - ): - results = await repo._search_hybrid(**HYBRID_KWARGS) + results = await fuse([FakeRow(id=1, score=5.0, title="fts-hit", content_snippet=content)], []) assert len(results) == 1 assert results[0].matched_chunk_text is None @@ -325,20 +234,7 @@ async def test_fts_only_result_does_not_copy_content_into_matched_chunk(): @pytest.mark.asyncio async def test_fts_only_result_with_null_content_keeps_null_matched_chunk(): """FTS-only results with no content_snippet should keep matched_chunk_text as None.""" - repo = ConcreteSearchRepo() - - fts_results = [ - FakeRow(id=1, score=5.0, title="fts-hit", content_snippet=None), - ] - vector_results = [] - - with ( - patch.object(repo, "search", new_callable=AsyncMock, return_value=fts_results), - patch.object( - repo, "_search_vector_only", new_callable=AsyncMock, return_value=vector_results - ), - ): - results = await repo._search_hybrid(**HYBRID_KWARGS) + results = await fuse([FakeRow(id=1, score=5.0, title="fts-hit", content_snippet=None)], []) assert len(results) == 1 assert results[0].matched_chunk_text is None @@ -347,8 +243,6 @@ async def test_fts_only_result_with_null_content_keeps_null_matched_chunk(): @pytest.mark.asyncio async def test_dual_source_result_keeps_vector_matched_chunk(): """Dual-source results should keep matched_chunk_text from vector search, not overwrite.""" - repo = ConcreteSearchRepo() - content = "Full note content from FTS." vector_chunk = "Specific chunk matched by vector search." fts_results = [ @@ -364,14 +258,22 @@ async def test_dual_source_result_keeps_vector_matched_chunk(): ), ] - with ( - patch.object(repo, "search", new_callable=AsyncMock, return_value=fts_results), - patch.object( - repo, "_search_vector_only", new_callable=AsyncMock, return_value=vector_results - ), - ): - results = await repo._search_hybrid(**HYBRID_KWARGS) + results = await fuse(fts_results, vector_results) assert len(results) == 1 - # Vector result overwrites the FTS row in rows_by_id, so matched_chunk_text is preserved + # Vector result overwrites the FTS row in rows_by_key, so matched_chunk_text is preserved assert results[0].matched_chunk_text == vector_chunk + + +@pytest.mark.asyncio +async def test_hybrid_fts_leg_runs_in_fts_mode_with_relaxation(): + """The lexical leg is the same prepared query in FTS mode, allowed to relax.""" + fts = FakeFts([FakeRow(id=1, score=5.0)]) + semantic = SemanticSearch(cast(Any, None), ProjectScope.single(1), fts, fake_vector_retrieval()) + + with patch.object(semantic, "vector_only", new_callable=AsyncMock, return_value=[]): + await semantic.hybrid(HYBRID_QUERY, limit=10, offset=0) + + assert [query.retrieval_mode for query in fts.queries] == [SearchRetrievalMode.FTS] + assert fts.queries[0].search_text == "test" + assert fts.calls[0]["allow_relaxed"] is True diff --git a/tests/repository/test_postgres_search_repository_unit.py b/tests/repository/test_postgres_search_repository_unit.py index 7ab175058..30a468e28 100644 --- a/tests/repository/test_postgres_search_repository_unit.py +++ b/tests/repository/test_postgres_search_repository_unit.py @@ -262,7 +262,7 @@ async def test_returns_empty_for_empty_embedding(self): embedding_provider=StubEmbeddingProvider(), ) session = AsyncMock() - result = await repo._run_vector_query(session, [], 10) + result = await repo._semantic_search()._run_vector_query(session, [], 10) assert result == [] diff --git a/tests/repository/test_rerank_pipeline.py b/tests/repository/test_rerank_pipeline.py index 97ca99783..3b1596325 100644 --- a/tests/repository/test_rerank_pipeline.py +++ b/tests/repository/test_rerank_pipeline.py @@ -1,7 +1,7 @@ """Rerank stage wiring in the shared search pipeline (vector + hybrid).""" from datetime import datetime, timezone -from typing import Any +from typing import Any, cast from unittest.mock import MagicMock import pytest @@ -19,16 +19,28 @@ demote_tail_scores, validate_rerank_scores, ) +from basic_memory.repository.search_filters import FtsBackend +from basic_memory.repository.search_query import PreparedSearchQuery +from basic_memory.repository.search_reader import ( + RERANK_POOL_CHUNK_FANOUT, + HydratedChunk, + Reranking, + SemanticSearch, + VectorRetrieval, + rerank_document_text, +) from basic_memory.repository.search_repository import create_search_repository -from basic_memory.repository.search_repository_base import RERANK_POOL_CHUNK_FANOUT +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_errors import ( RerankProviderContractError, RerankTransientError, SemanticDependenciesMissingError, ) +from basic_memory.repository.semantic_vector_index import SemanticVectorIndex from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode from basic_memory.services.entity_service import EntityService +from tests.repository.test_hybrid_fusion import FakeFts type BackendSearchRepository = SQLiteSearchRepository | PostgresSearchRepository @@ -191,69 +203,63 @@ def _row(**overrides) -> SearchIndexRow: return SearchIndexRow(**base) -def _unit_repo() -> SQLiteSearchRepository: - """A repo built without a real DB — for the pure rerank helper methods.""" - config = BasicMemoryConfig( - env="test", - projects={"test-project": "/tmp/test"}, - default_project="test-project", - database_backend=DatabaseBackend.SQLITE, - semantic_search_enabled=True, +def _reranking(provider: Any, *, candidates: int = 20, max_document_chars: int = 0) -> Reranking: + return Reranking( + provider=provider, candidates=candidates, max_document_chars=max_document_chars ) - return SQLiteSearchRepository( - MagicMock(), - project_id=1, - app_config=config, + + +def _semantic( + *, + vector_k: int = 100, + rerank: Reranking | None = None, + fts: FtsBackend | None = None, +) -> SemanticSearch: + """A pipeline built without a real DB, for the pure rerank helpers.""" + vector = VectorRetrieval( + index=cast(SemanticVectorIndex, object()), + index_name="sqlite-vec", embedding_provider=_StubEmbeddingProvider(), + embedding_model="stub:4", + vector_k=vector_k, + min_similarity=0.0, ) + return SemanticSearch(MagicMock(), ProjectScope.single(1), fts or FakeFts(), vector, rerank) # --- Pure helper behavior --- -def test_should_rerank_gating(): - repo = _unit_repo() - repo._rerank_provider = None - assert repo._should_rerank("auth") is False - repo._rerank_provider = _FakeReranker({}) - assert repo._should_rerank("") is False - assert repo._should_rerank("auth") is True +def test_rerank_is_active_only_with_a_provider_and_a_query(): + assert _semantic()._active_rerank("auth") is None + semantic = _semantic(rerank=_reranking(_FakeReranker({}))) + assert semantic._active_rerank("") is None + assert semantic._active_rerank("auth") is semantic.rerank def test_rerank_document_text_fallbacks(): - repo = _unit_repo() - assert repo._rerank_document_text(_row(title="T", matched_chunk_text="chunk")) == "chunk\nT" + assert rerank_document_text(_row(title="T", matched_chunk_text="chunk"), 0) == "chunk\nT" assert ( - repo._rerank_document_text(_row(title="T", matched_chunk_text=None, content_snippet="snip")) + rerank_document_text(_row(title="T", matched_chunk_text=None, content_snippet="snip"), 0) == "snip\nT" ) - assert ( - repo._rerank_document_text(_row(title=None, matched_chunk_text="only-body")) == "only-body" - ) - assert ( - repo._rerank_document_text(_row(title="only-title", content_snippet=None)) == "only-title" - ) - assert repo._rerank_document_text(_row(title=None, content_snippet=None)) == "" + assert rerank_document_text(_row(title=None, matched_chunk_text="only-body"), 0) == "only-body" + assert rerank_document_text(_row(title="only-title", content_snippet=None), 0) == "only-title" + assert rerank_document_text(_row(title=None, content_snippet=None), 0) == "" def test_rerank_document_text_truncation(): - repo = _unit_repo() row = _row(title="T", matched_chunk_text="x" * 500) # full text = 500 + "\nT" = 502 chars - repo._reranker_max_document_chars = 0 # disabled - assert len(repo._rerank_document_text(row)) == 502 - repo._reranker_max_document_chars = 100 # trims to the leading (most-relevant) text - trimmed = repo._rerank_document_text(row) + assert len(rerank_document_text(row, 0)) == 502 # disabled + trimmed = rerank_document_text(row, 100) # trims to the leading (most-relevant) text assert len(trimmed) == 100 and trimmed == "x" * 100 - repo._reranker_max_document_chars = 10_000 # no-op when already under the cap - assert len(repo._rerank_document_text(row)) == 502 + assert len(rerank_document_text(row, 10_000)) == 502 # no-op when already under the cap def test_rerank_document_text_cap_preserves_matched_body_with_long_title(): - repo = _unit_repo() - repo._reranker_max_document_chars = 8 row = _row(title="title-" * 20, matched_chunk_text="MATCHED body") - assert "MATCHED" in repo._rerank_document_text(row) + assert "MATCHED" in rerank_document_text(row, 8) def test_demoted_tail_scores_are_stable_as_the_tail_grows(): @@ -268,66 +274,58 @@ def test_demoted_tail_scores_are_stable_as_the_tail_grows(): def test_candidate_limit_over_fetches_chunks_for_rerank_pool(): """With reranking active, over-fetch chunks so dedup can't starve the rerank window.""" - repo = _unit_repo() - repo._semantic_vector_k = 5 - repo._reranker_candidates = 20 + plain = _semantic(vector_k=5) + assert plain._candidate_limit(limit=1, offset=0, query_text="auth") == 10 # max(5, 10) - repo._rerank_provider = None - assert repo._candidate_limit(limit=1, offset=0, query_text="auth") == 10 # max(5, 10) - - repo._rerank_provider = _FakeReranker({}) + reranked = _semantic(vector_k=5, rerank=_reranking(_FakeReranker({}), candidates=20)) assert ( - repo._candidate_limit(limit=1, offset=0, query_text="auth") == 20 * RERANK_POOL_CHUNK_FANOUT + reranked._candidate_limit(limit=1, offset=0, query_text="auth") + == 20 * RERANK_POOL_CHUNK_FANOUT ) - assert repo._candidate_limit(limit=1, offset=0, query_text="") == 10 # no query → no bump + assert reranked._candidate_limit(limit=1, offset=0, query_text="") == 10 # no query → no bump def test_candidate_limit_expands_only_for_results_beyond_rerank_pool(): """The fixed rerank window grows only enough to supply the requested tail.""" - repo = _unit_repo() - repo._semantic_vector_k = 5 - repo._reranker_candidates = 20 - repo._rerank_provider = _FakeReranker({}) + semantic = _semantic(vector_k=5, rerank=_reranking(_FakeReranker({}), candidates=20)) - first_page_limit = repo._candidate_limit(limit=10, offset=0, query_text="auth") + first_page_limit = semantic._candidate_limit(limit=10, offset=0, query_text="auth") assert first_page_limit == 20 * RERANK_POOL_CHUNK_FANOUT - assert repo._candidate_limit(limit=20, offset=0, query_text="auth") == first_page_limit - assert repo._candidate_limit(limit=10, offset=10, query_text="auth") == first_page_limit - assert repo._candidate_limit(limit=21, offset=0, query_text="auth") == 90 - assert repo._candidate_limit(limit=10, offset=19, query_text="auth") == 170 - assert repo._candidate_limit(limit=10, offset=20, query_text="auth") == 180 + assert semantic._candidate_limit(limit=20, offset=0, query_text="auth") == first_page_limit + assert semantic._candidate_limit(limit=10, offset=10, query_text="auth") == first_page_limit + assert semantic._candidate_limit(limit=21, offset=0, query_text="auth") == 90 + assert semantic._candidate_limit(limit=10, offset=19, query_text="auth") == 170 + assert semantic._candidate_limit(limit=10, offset=20, query_text="auth") == 180 # Large first pages still retrieve their untouched tail and pagination probe. - assert repo._candidate_limit(limit=101, offset=0, query_text="auth") == 890 + assert semantic._candidate_limit(limit=101, offset=0, query_text="auth") == 890 @pytest.mark.asyncio async def test_rerank_paginate_noop_paths(): - repo = _unit_repo() rows = [_row(id=1), _row(id=2)] - repo._rerank_provider = None - assert await repo._rerank_and_paginate("auth", rows, offset=0, limit=10) == rows + assert await _semantic()._rerank_and_paginate("auth", rows, offset=0, limit=10) == rows reranker = _FakeReranker({}) - repo._rerank_provider = reranker - assert await repo._rerank_and_paginate("", rows, offset=0, limit=10) == rows - assert await repo._rerank_and_paginate("auth", [], offset=0, limit=10) == [] - assert await repo._rerank_and_paginate("auth", rows, offset=2, limit=10) == [] + semantic = _semantic(rerank=_reranking(reranker)) + assert await semantic._rerank_and_paginate("", rows, offset=0, limit=10) == rows + assert await semantic._rerank_and_paginate("auth", [], offset=0, limit=10) == [] + assert await semantic._rerank_and_paginate("auth", rows, offset=2, limit=10) == [] assert reranker.calls == 0 @pytest.mark.asyncio async def test_rerank_paginate_reorders_rescore_and_demotes_tail(): - repo = _unit_repo() - repo._rerank_provider = _FakeReranker({"Alpha": 0.1, "Bravo": 0.9, "Charlie": 0.5}) - repo._reranker_candidates = 2 + semantic = _semantic( + rerank=_reranking(_FakeReranker({"Alpha": 0.1, "Bravo": 0.9, "Charlie": 0.5}), candidates=2) + ) rows = [ _row(id=1, title="Alpha"), _row(id=2, title="Bravo"), _row(id=3, title="Charlie"), # past the pool ] - result = await repo._rerank_and_paginate("auth", rows, offset=0, limit=3) + result = await semantic._rerank_and_paginate("auth", rows, offset=0, limit=3) assert [r.title for r in result] == ["Bravo", "Alpha", "Charlie"] assert result[0].score == 0.9 # reranker relevance replaces the prior score @@ -343,16 +341,16 @@ async def test_rerank_paginate_reorders_rescore_and_demotes_tail(): @pytest.mark.asyncio async def test_rerank_paginate_preserves_pool_before_tail_at_zero_floor(): - repo = _unit_repo() - repo._rerank_provider = _FakeReranker({"Alpha": 0.0, "Bravo": 0.9}) - repo._reranker_candidates = 2 + semantic = _semantic( + rerank=_reranking(_FakeReranker({"Alpha": 0.0, "Bravo": 0.9}), candidates=2) + ) rows = [ _row(id=1, title="Alpha"), _row(id=2, title="Bravo"), _row(id=3, title="Charlie"), ] - result = await repo._rerank_and_paginate("auth", rows, offset=0, limit=3) + result = await semantic._rerank_and_paginate("auth", rows, offset=0, limit=3) assert [row.title for row in result] == ["Bravo", "Alpha", "Charlie"] assert [row.score for row in result] == [0.9, 0.0, 0.0] @@ -361,17 +359,15 @@ async def test_rerank_paginate_preserves_pool_before_tail_at_zero_floor(): @pytest.mark.asyncio async def test_rerank_paginate_scores_singleton_prefix_and_demotes_tail(): """A one-row prefix still calibrates scores before cross-project merging.""" - repo = _unit_repo() reranker = _FakeReranker({"Only": 0.4}) - repo._rerank_provider = reranker - repo._reranker_candidates = 2 + semantic = _semantic(rerank=_reranking(reranker, candidates=2)) stable_rows = [_row(id=1, title="Only", score=0.5)] expanded_rows = [ stable_rows[0], _row(id=2, title="Tail", score=1.3), ] - result = await repo._rerank_and_paginate( + result = await semantic._rerank_and_paginate( "auth", expanded_rows, offset=0, @@ -388,10 +384,8 @@ async def test_rerank_paginate_scores_singleton_prefix_and_demotes_tail(): @pytest.mark.asyncio async def test_rerank_paginate_calibrates_tail_scores_on_deep_page(): """Deep pages rescore the fixed prefix before returning its calibrated tail.""" - repo = _unit_repo() reranker = _FakeReranker({"n1": 0.9, "n2": 0.8}) - repo._rerank_provider = reranker - repo._reranker_candidates = 2 + semantic = _semantic(rerank=_reranking(reranker, candidates=2)) stable_rows = [_row(id=1, title="n1"), _row(id=2, title="n2")] expanded_rows = [ _row(id=3, title="newly strengthened"), @@ -401,7 +395,7 @@ async def test_rerank_paginate_calibrates_tail_scores_on_deep_page(): _row(id=5, title="n5"), ] - result = await repo._rerank_and_paginate( + result = await semantic._rerank_and_paginate( "auth", expanded_rows, offset=2, @@ -417,10 +411,8 @@ async def test_rerank_paginate_calibrates_tail_scores_on_deep_page(): @pytest.mark.asyncio async def test_rerank_paginate_keeps_expanded_candidates_out_of_stable_prefix(): """A larger tail retrieval cannot replace candidates in the reranked prefix.""" - repo = _unit_repo() reranker = _FakeReranker({"Alpha": 0.1, "Bravo": 0.9, "Charlie": 1.0}) - repo._rerank_provider = reranker - repo._reranker_candidates = 2 + semantic = _semantic(rerank=_reranking(reranker, candidates=2)) stable_rows = [_row(id=1, title="Alpha"), _row(id=2, title="Bravo")] expanded_rows = [ _row(id=3, title="Charlie"), @@ -429,14 +421,14 @@ async def test_rerank_paginate_keeps_expanded_candidates_out_of_stable_prefix(): _row(id=4, title="Delta"), ] - result = await repo._rerank_and_paginate( + result = await semantic._rerank_and_paginate( "auth", expanded_rows, offset=0, limit=3, stable_rows=stable_rows, ) - expanded_result = await repo._rerank_and_paginate( + expanded_result = await semantic._rerank_and_paginate( "auth", expanded_rows + [_row(id=5, title="Echo"), _row(id=6, title="Foxtrot")], offset=0, @@ -453,38 +445,32 @@ async def test_rerank_paginate_keeps_expanded_candidates_out_of_stable_prefix(): @pytest.mark.asyncio async def test_rerank_paginate_surfaces_transient_provider_error(): """Transient failures must not silently replace reranked order with retrieval order.""" - repo = _unit_repo() - repo._rerank_provider = _ExplodingReranker() - repo._reranker_candidates = 20 + semantic = _semantic(rerank=_reranking(_ExplodingReranker(), candidates=20)) rows = [_row(id=1, title="A"), _row(id=2, title="B")] with pytest.raises(RerankTransientError, match="backend unreachable"): - await repo._rerank_and_paginate("auth", rows, offset=0, limit=10) + await semantic._rerank_and_paginate("auth", rows, offset=0, limit=10) @pytest.mark.asyncio async def test_rerank_paginate_does_not_duplicate_results_when_later_page_is_transient(): """A later page fails instead of changing order and repeating an earlier result.""" - repo = _unit_repo() - repo._rerank_provider = _SucceedsThenTransientReranker() - repo._reranker_candidates = 2 + semantic = _semantic(rerank=_reranking(_SucceedsThenTransientReranker(), candidates=2)) rows = [_row(id=1, title="A"), _row(id=2, title="B")] - first_page = await repo._rerank_and_paginate("auth", rows, offset=0, limit=1) + first_page = await semantic._rerank_and_paginate("auth", rows, offset=0, limit=1) assert [row.id for row in first_page] == [2] with pytest.raises(RerankTransientError, match="backend unreachable"): - await repo._rerank_and_paginate("auth", rows, offset=1, limit=1) + await semantic._rerank_and_paginate("auth", rows, offset=1, limit=1) @pytest.mark.asyncio async def test_rerank_paginate_misaligned_scores_raise(): """A length mismatch is a provider bug — fail fast, don't degrade.""" - repo = _unit_repo() - repo._rerank_provider = _BadReranker() - repo._reranker_candidates = 20 + semantic = _semantic(rerank=_reranking(_BadReranker(), candidates=20)) with pytest.raises(RerankProviderContractError, match="Reranker returned 0 scores"): - await repo._rerank_and_paginate("auth", [_row(id=1), _row(id=2)], offset=0, limit=10) + await semantic._rerank_and_paginate("auth", [_row(id=1), _row(id=2)], offset=0, limit=10) @pytest.mark.asyncio @@ -498,11 +484,9 @@ async def test_rerank_paginate_misaligned_scores_raise(): ) async def test_rerank_paginate_surfaces_permanent_faults(exc): """Permanent faults (contract break, missing deps) propagate — not silently degraded.""" - repo = _unit_repo() - repo._rerank_provider = _PermanentFaultReranker(exc) - repo._reranker_candidates = 20 + semantic = _semantic(rerank=_reranking(_PermanentFaultReranker(exc), candidates=20)) with pytest.raises(type(exc)): - await repo._rerank_and_paginate("auth", [_row(id=1), _row(id=2)], offset=0, limit=10) + await semantic._rerank_and_paginate("auth", [_row(id=1), _row(id=2)], offset=0, limit=10) # --- End-to-end through both repository backends --- @@ -601,17 +585,20 @@ async def test_vector_search_expands_tail_from_stable_rerank_pool( rerank_search_repository._rerank_provider = _FakeReranker({"Alpha": 0.1, "Bravo": 0.9}) candidate_limits: list[int] = [] - run_vector_query = rerank_search_repository._run_vector_query + run_vector_query = SemanticSearch._run_vector_query async def record_vector_query( + self: SemanticSearch, session: Any, query_embedding: list[float], candidate_limit: int, - ) -> list[dict[str, Any]]: + *, + trace: Any = None, + ) -> list[HydratedChunk]: candidate_limits.append(candidate_limit) - return await run_vector_query(session, query_embedding, candidate_limit) + return await run_vector_query(self, session, query_embedding, candidate_limit, trace=trace) - monkeypatch.setattr(rerank_search_repository, "_run_vector_query", record_vector_query) + monkeypatch.setattr(SemanticSearch, "_run_vector_query", record_vector_query) results = await rerank_search_repository.search( search_text="auth session token", @@ -637,12 +624,12 @@ async def slow_rerank(query: str, documents: list[str]) -> list[float]: rerank_search_repository._rerank_provider = reranker monkeypatch.setattr(reranker, "rerank", slow_rerank) monkeypatch.setattr( - "basic_memory.repository.search_repository_base.time.perf_counter", + "basic_memory.repository.search_reader.time.perf_counter", lambda: clock["now"], ) warning = MagicMock() monkeypatch.setattr( - "basic_memory.repository.search_repository_base.logger.warning", + "basic_memory.repository.search_reader.logger.warning", warning, ) @@ -751,17 +738,20 @@ async def test_hybrid_search_preserves_candidate_windows( rerank_search_repository._rerank_provider = None candidate_limits: list[int] = [] - run_vector_query = rerank_search_repository._run_vector_query + run_vector_query = SemanticSearch._run_vector_query async def record_vector_query( + self: SemanticSearch, session: Any, query_embedding: list[float], candidate_limit: int, - ) -> list[dict[str, Any]]: + *, + trace: Any = None, + ) -> list[HydratedChunk]: candidate_limits.append(candidate_limit) - return await run_vector_query(session, query_embedding, candidate_limit) + return await run_vector_query(self, session, query_embedding, candidate_limit, trace=trace) - monkeypatch.setattr(rerank_search_repository, "_run_vector_query", record_vector_query) + monkeypatch.setattr(SemanticSearch, "_run_vector_query", record_vector_query) baseline_results = await rerank_search_repository.search( search_text="auth session token", @@ -801,17 +791,12 @@ async def record_vector_query( @pytest.mark.asyncio async def test_hybrid_search_keeps_deep_tail_stable_as_candidate_window_grows(monkeypatch): """Late dual-source evidence cannot move a row across an earlier tail page.""" - repo = _unit_repo() - repo._semantic_vector_k = 2 - repo._reranker_candidates = 2 reranker = _FakeReranker({"Alpha": 0.9, "Bravo": 0.8}) - repo._rerank_provider = reranker charlie = _row(id=3, title="Charlie") delta = _row(id=4, title="Delta") - async def fake_fts_search(*args, limit: int, **kwargs) -> list[SearchIndexRow]: - assert kwargs["retrieval_mode"] == SearchRetrievalMode.FTS + def fts_window(limit: int) -> list[SearchIndexRow]: if limit <= 8: return [ _row(id=1, title="Alpha", score=10.0), @@ -824,8 +809,12 @@ async def fake_fts_search(*args, limit: int, **kwargs) -> list[SearchIndexRow]: _row(id=4, title="Delta", score=7.0), ] - async def fake_vector_search(**kwargs) -> list[SearchIndexRow]: + fts = FakeFts(fts_window) + semantic = _semantic(vector_k=2, rerank=_reranking(reranker, candidates=2), fts=fts) + + async def fake_vector_search(query: PreparedSearchQuery, **kwargs: Any) -> list[SearchIndexRow]: candidate_limit = kwargs["candidate_limit"] + assert candidate_limit is not None if candidate_limit <= 8: return [ _row(id=1, title="Alpha", score=1.0), @@ -844,22 +833,11 @@ async def fake_vector_search(**kwargs) -> list[SearchIndexRow]: _row(id=4, title="Delta", score=0.7), ] - monkeypatch.setattr(repo, "search", fake_fts_search) - monkeypatch.setattr(repo, "_search_vector_only", fake_vector_search) + monkeypatch.setattr(semantic, "vector_only", fake_vector_search) async def deep_page(offset: int) -> list[SearchIndexRow]: - return await repo._search_hybrid( - search_text="auth", - permalink=None, - permalink_match=None, - title=None, - note_types=None, - after_date=None, - search_item_types=None, - categories=None, - metadata_filters=None, - file_path_prefix=None, - temporal=None, + return await semantic.hybrid( + PreparedSearchQuery(search_text="auth", retrieval_mode=SearchRetrievalMode.HYBRID), limit=1, offset=offset, ) @@ -870,6 +848,7 @@ async def deep_page(offset: int) -> list[SearchIndexRow]: assert [row.id for row in first_tail_page] == [charlie.id] assert [row.id for row in second_tail_page] == [delta.id] assert reranker.calls == 2 + assert all(query.retrieval_mode == SearchRetrievalMode.FTS for query in fts.queries) @pytest.mark.asyncio diff --git a/tests/repository/test_search_file_path_prefix.py b/tests/repository/test_search_file_path_prefix.py index afe4be233..b668218d2 100644 --- a/tests/repository/test_search_file_path_prefix.py +++ b/tests/repository/test_search_file_path_prefix.py @@ -10,6 +10,7 @@ from datetime import datetime, timezone from types import SimpleNamespace +from typing import cast from unittest.mock import AsyncMock import pytest @@ -18,6 +19,8 @@ from basic_memory.models.knowledge import Entity from basic_memory.repository.search_index_row import SearchIndexRow from basic_memory.repository.search_filters import file_path_prefix_condition +from basic_memory.repository.search_reader import HydratedChunk, SemanticSearch +from basic_memory.repository.semantic_vector_index import SemanticVectorIndex from basic_memory.schemas.search import ( SearchItemType, SearchRetrievalMode, @@ -305,18 +308,24 @@ async def test_semantic_retrieval_honors_the_scope( ), ) monkeypatch.setattr(search_repository, "_ensure_vector_tables", AsyncMock()) - monkeypatch.setattr(search_repository, "_prepare_vector_session", AsyncMock()) + # The nearest-neighbour stage is stubbed, so the adapter is never consulted. monkeypatch.setattr( search_repository, + "_semantic_vector_index", + cast(SemanticVectorIndex, object()), + raising=False, + ) + monkeypatch.setattr( + SemanticSearch, "_run_vector_query", AsyncMock( return_value=[ - { - "entity_id": row_id, - "chunk_key": f"entity:{row_id}:0", - "chunk_text": "subtree scope fixture", - "best_similarity": 0.9, - } + HydratedChunk( + entity_id=row_id, + chunk_key=f"entity:{row_id}:0", + chunk_text="subtree scope fixture", + similarity=0.9, + ) for row_id in seeded_paths.values() ] ), diff --git a/tests/repository/test_search_reader.py b/tests/repository/test_search_reader.py new file mode 100644 index 000000000..ed6da6e62 --- /dev/null +++ b/tests/repository/test_search_reader.py @@ -0,0 +1,283 @@ +"""SearchReader dispatch, and the SemanticSearch edges the pipeline tests do not reach.""" + +from dataclasses import replace +from itertools import count +from typing import Any, cast +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from basic_memory.repository.search_query import PreparedSearchQuery +from basic_memory.repository.search_reader import ( + HydratedChunk, + Reranking, + SearchReader, + SemanticSearch, + parse_chunk_key, + vector_eligible, +) +from basic_memory.repository.search_scope import ProjectScope +from basic_memory.repository.search_trace import ManifestReadiness, SearchTraceCollector +from basic_memory.repository.semantic_errors import SemanticSearchDisabledError +from basic_memory.schemas.search import SearchRetrievalMode +from tests.repository.test_hybrid_fusion import ( + HYBRID_QUERY, + FakeFts, + FakeRow, + fake_vector_retrieval, +) +from tests.repository.test_vector_threshold import run_vector_only, vector_semantic + +SCOPE = ProjectScope.single(1) +SEMANTIC_MODES = [SearchRetrievalMode.VECTOR, SearchRetrievalMode.HYBRID] + + +def _semantic(fts: FakeFts | None = None) -> SemanticSearch: + return SemanticSearch(cast(Any, None), SCOPE, fts or FakeFts(), fake_vector_retrieval()) + + +# --- Eligibility and keys --- + + +@pytest.mark.parametrize( + ("query", "eligible"), + [ + (PreparedSearchQuery(search_text="auth"), True), + (PreparedSearchQuery(search_text=" auth "), True), + (PreparedSearchQuery(search_text=None), False), + (PreparedSearchQuery(search_text=" "), False), + (PreparedSearchQuery(search_text="*"), False), + (PreparedSearchQuery(search_text="auth", permalink="specs/auth"), False), + (PreparedSearchQuery(search_text="auth", permalink_match="specs/*"), False), + (PreparedSearchQuery(search_text="auth", title="Auth"), False), + ], +) +def test_vector_eligible_requires_text_and_no_identity_filter(query, eligible): + assert vector_eligible(query) is eligible + + +def test_parse_chunk_key_reads_type_and_row_id(): + assert parse_chunk_key("observation:5:0") == ("observation", 5) + with pytest.raises(ValueError): + parse_chunk_key("entity:not-a-number:0") + with pytest.raises(IndexError): + parse_chunk_key("garbage") + + +# --- SearchReader dispatch --- + + +@pytest.mark.asyncio +async def test_fts_mode_hands_the_engine_the_query_and_every_option(): + fts = FakeFts([FakeRow(id=1)]) + reader = SearchReader(SCOPE, fts) + query = PreparedSearchQuery(search_text="auth") + + rows = await reader.search( + query, + limit=5, + offset=2, + allow_relaxed=True, + candidate_keys=[("entity", 1)], + ) + + assert [row.id for row in rows] == [1] + assert fts.queries == [query] + assert fts.calls == [ + {"limit": 5, "offset": 2, "allow_relaxed": True, "candidate_keys": [("entity", 1)]} + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", SEMANTIC_MODES) +async def test_semantic_modes_reject_queries_with_nothing_to_embed(mode): + reader = SearchReader(SCOPE, FakeFts(), _semantic()) + + with pytest.raises(ValueError, match="requires a non-empty text query"): + await reader.search( + PreparedSearchQuery(title="Auth", retrieval_mode=mode), limit=10, offset=0 + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", SEMANTIC_MODES) +async def test_semantic_modes_without_a_semantic_stack_raise_disabled(mode): + """A reader built without the semantic stack answers full-text queries only.""" + reader = SearchReader(SCOPE, FakeFts()) + + with pytest.raises(SemanticSearchDisabledError): + await reader.search( + PreparedSearchQuery(search_text="auth", retrieval_mode=mode), limit=10, offset=0 + ) + + +@pytest.mark.asyncio +async def test_semantic_modes_route_to_the_pipeline(): + semantic = _semantic() + reader = SearchReader(SCOPE, FakeFts(), semantic) + vector_query = PreparedSearchQuery( + search_text="auth", retrieval_mode=SearchRetrievalMode.VECTOR + ) + + with patch.object( + semantic, "vector_only", new_callable=AsyncMock, return_value=[FakeRow(id=1)] + ) as vector_only: + rows = await reader.search(vector_query, limit=3, offset=1) + assert [row.id for row in rows] == [1] + vector_only.assert_awaited_once_with(vector_query, limit=3, offset=1, trace=None) + + with patch.object( + semantic, "hybrid", new_callable=AsyncMock, return_value=[FakeRow(id=2)] + ) as hybrid: + rows = await reader.search(HYBRID_QUERY, limit=3, offset=1) + assert [row.id for row in rows] == [2] + hybrid.assert_awaited_once_with(HYBRID_QUERY, limit=3, offset=1, trace=None) + + +@pytest.mark.asyncio +async def test_count_is_full_text_only(): + reader = SearchReader(SCOPE, FakeFts([FakeRow(id=1), FakeRow(id=2)])) + + assert await reader.count(PreparedSearchQuery(search_text="auth"), allow_relaxed=True) == 2 + with pytest.raises(ValueError, match="Exact counts are only supported"): + await reader.count( + PreparedSearchQuery(search_text="auth", retrieval_mode=SearchRetrievalMode.VECTOR) + ) + + +# --- SemanticSearch edges --- + + +@pytest.mark.asyncio +async def test_empty_candidate_window_skips_the_adapter(): + adapter = MagicMock() + adapter.search = AsyncMock(side_effect=AssertionError("adapter must not be consulted")) + vector = replace(fake_vector_retrieval(), index=adapter) + semantic = SemanticSearch(cast(Any, None), SCOPE, FakeFts(), vector) + + assert await semantic._run_vector_query(AsyncMock(), [0.1], 0) == [] + + +@pytest.mark.asyncio +async def test_no_row_ids_means_no_row_fetch(): + session_maker = MagicMock(side_effect=AssertionError("no session for an empty fetch")) + semantic = SemanticSearch(session_maker, SCOPE, FakeFts(), fake_vector_retrieval()) + + assert await semantic._fetch_search_index_rows_by_ids([]) == {} + + +@pytest.mark.asyncio +async def test_vector_only_with_no_parseable_chunk_keys_returns_nothing(): + semantic = vector_semantic() + fetch_rows = AsyncMock() + rows = [HydratedChunk(entity_id=0, chunk_key="garbage", chunk_text="bad", similarity=0.95)] + + assert await run_vector_only(semantic, rows, fetch_rows) == [] + fetch_rows.assert_not_called() + + +@pytest.mark.asyncio +async def test_hybrid_skips_rows_without_an_id_on_both_legs(): + fts = FakeFts([FakeRow(id=None, score=2.0), FakeRow(id=1, score=5.0)]) + semantic = _semantic(fts) + vector_results = [FakeRow(id=None, score=0.9), FakeRow(id=2, score=0.8)] + + with patch.object(semantic, "vector_only", new_callable=AsyncMock, return_value=vector_results): + rows = await semantic.hybrid(HYBRID_QUERY, limit=10, offset=0) + + assert [(row.type, row.id) for row in rows] == [("entity", 1), ("entity", 2)] + + +@pytest.mark.asyncio +async def test_hybrid_slow_query_warning_names_the_scope(monkeypatch): + semantic = _semantic(FakeFts([FakeRow(id=1, score=5.0)])) + clock = count(0.0, 3.0) + monkeypatch.setattr( + "basic_memory.repository.search_reader.time.perf_counter", lambda: next(clock) + ) + warning = MagicMock() + monkeypatch.setattr("basic_memory.repository.search_reader.logger.warning", warning) + + with patch.object(semantic, "vector_only", new_callable=AsyncMock, return_value=[]): + await semantic.hybrid(HYBRID_QUERY, limit=10, offset=0) + + warning.assert_called_once() + assert warning.call_args.args[0].startswith("[SEMANTIC_SLOW_QUERY]") + assert warning.call_args.kwargs["retrieval_mode"] == "hybrid" + assert warning.call_args.kwargs["scope"] == (1,) + + +@pytest.mark.asyncio +async def test_built_in_adapter_reads_manifest_readiness_when_tracing(monkeypatch): + """A traced built-in lookup reports how much of the manifest retrieval can answer from.""" + adapter = MagicMock() + adapter.search = AsyncMock(return_value=[]) + semantic = SemanticSearch( + cast(Any, None), SCOPE, FakeFts(), replace(fake_vector_retrieval(), index=adapter) + ) + readiness = ManifestReadiness( + configured_index="sqlite-vec", + configured_model="fake:384", + ready_rows=3, + pending_rows=1, + other_identity_rows=0, + ) + read_readiness = AsyncMock(return_value=readiness) + monkeypatch.setattr( + "basic_memory.repository.search_reader.read_manifest_readiness", read_readiness + ) + trace = SearchTraceCollector() + session = AsyncMock() + + assert await semantic._run_vector_query(session, [0.1], 5, trace=trace) == [] + + assert trace.readiness is readiness + read_readiness.assert_awaited_once_with(session, SCOPE, "sqlite-vec", "fake:384") + adapter.search.assert_awaited_once_with([0.1], limit=5) + + +@pytest.mark.asyncio +async def test_fts_gate_zeroes_weak_lexical_scores(monkeypatch): + """Below the gate a lexical hit contributes nothing to fusion instead of a sliver.""" + monkeypatch.setattr("basic_memory.repository.search_reader.FTS_GATE_THRESHOLD", 0.5) + semantic = _semantic(FakeFts([FakeRow(id=1, score=10.0), FakeRow(id=2, score=1.0)])) + + with patch.object(semantic, "vector_only", new_callable=AsyncMock, return_value=[]): + rows = await semantic.hybrid(HYBRID_QUERY, limit=10, offset=0) + + assert [(row.id, row.score) for row in rows] == [(1, 1.0), (2, 0.0)] + + +class _PrefixReranker: + """Scores documents by position so the rerank stage is deterministic.""" + + model_name = "prefix" + + async def rerank(self, query: str, documents: list[str]) -> list[float]: + return [1.0 - index * 0.1 for index in range(len(documents))] + + def runtime_log_attrs(self) -> dict[str, Any]: + return {} + + +@pytest.mark.asyncio +async def test_hybrid_trace_records_the_stable_pool_refetch(): + """A page past the fixed rerank prefix refetches the stable pool, and the trace says so.""" + rows = [ + FakeRow(id=index, score=10.0 - index, title=f"n{index}", entity_id=index) + for index in range(1, 6) + ] + semantic = SemanticSearch( + cast(Any, None), + SCOPE, + FakeFts(rows), + replace(fake_vector_retrieval(), vector_k=2), + Reranking(provider=_PrefixReranker(), candidates=2, max_document_chars=0), + ) + trace = SearchTraceCollector() + + with patch.object(semantic, "vector_only", new_callable=AsyncMock, return_value=[]): + page = await semantic.hybrid(HYBRID_QUERY, limit=1, offset=3, trace=trace) + + assert trace.stable_pool_refetched is True + assert [row.id for row in page] == [4] diff --git a/tests/repository/test_search_trace.py b/tests/repository/test_search_trace.py index c6b913bcb..89efb107d 100644 --- a/tests/repository/test_search_trace.py +++ b/tests/repository/test_search_trace.py @@ -14,7 +14,7 @@ from basic_memory.config import BasicMemoryConfig, DatabaseBackend from basic_memory.repository.postgres_search_repository import PostgresSearchRepository from basic_memory.repository.search_index_row import SearchIndexRow -from basic_memory.repository.search_repository_base import FUSION_BONUS +from basic_memory.repository.search_reader import FUSION_BONUS from basic_memory.repository.search_trace import ( BelowThreshold, FilteredOut, @@ -751,7 +751,7 @@ async def test_classify_hydration_drop_observes_pending_to_ready_transition( ) await session.commit() async with db.scoped_session(session_maker) as session: - assert await repository._hydrate_vector_matches(session, [match]) == [] + assert await repository._semantic_search()._hydrate_vector_matches(session, [match]) == [] async with db.scoped_session(session_maker) as session: await session.execute( diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index d9ae052f1..4a5857aca 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -8,16 +8,21 @@ from contextlib import asynccontextmanager from datetime import datetime from types import SimpleNamespace -from typing import override, Any +from typing import override, Any, cast from unittest.mock import AsyncMock, Mock import pytest import basic_memory.repository.search_repository_base as search_repository_base_module +from basic_memory.repository.embedding_provider import EmbeddingProvider from basic_memory.repository.fastembed_provider import FastEmbedEmbeddingProvider -from basic_memory.repository.search_index_row import SearchIndexRow +from basic_memory.repository.search_index_row import SearchIndexKey, SearchIndexRow +from basic_memory.repository.search_reader import ( + HydratedChunk, + SemanticSearch, + VectorRetrieval, +) from basic_memory.repository.search_repository_base import ( - SearchIndexKey, SearchRepositoryBase, _PreparedEntityVectorSync, ) @@ -37,6 +42,7 @@ from basic_memory.repository.semantic_vector_sync import PendingEmbeddingJob from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode from basic_memory.temporal import TemporalFilter +from tests.repository.test_hybrid_fusion import FakeFts # --- Helpers --- @@ -67,6 +73,7 @@ def __init__(self): self.session_maker = None self.project_id = 1 self.scope = ProjectScope.single(1) + self._fts = FakeFts() @override async def init_search_index(self): @@ -113,17 +120,6 @@ async def search( async def _ensure_vector_tables(self): pass - @override - async def _run_vector_query( - self, - session, - query_embedding, - candidate_limit, - *, - trace: SearchTraceCollector | None = None, - ): - return [] - @override async def _write_embeddings(self, session, jobs, embeddings): pass @@ -146,10 +142,6 @@ async def _delete_stale_chunks( async def _update_timestamp_sql(self): return "CURRENT_TIMESTAMP" - @override - def _distance_to_similarity(self, distance: float) -> float: - return 1.0 / (1.0 + max(distance, 0.0)) - class _RecordingVectorIndex: """Protocol-complete adapter that records generation-safe upserts.""" @@ -185,12 +177,25 @@ async def search( return [] +def _semantic_search(*, index_name: str, adapter: Any = None) -> SemanticSearch: + """The vector pipeline over an adapter the test controls.""" + vector = VectorRetrieval( + index=adapter if adapter is not None else _RecordingVectorIndex(), + index_name=index_name, + embedding_provider=cast( + EmbeddingProvider, SimpleNamespace(model_name="stub", dimensions=4) + ), + embedding_model="stub:4", + vector_k=100, + min_similarity=0.0, + ) + return SemanticSearch(cast(Any, None), ProjectScope.single(1), FakeFts(), vector) + + @pytest.mark.asyncio async def test_vector_match_hydration_batches_large_adapter_results() -> None: """Deep vector pages must not create an unbounded SQL bind-parameter list.""" - repo = _ConcreteRepo() - repo._semantic_vector_index_name = "milvus" - repo._embedding_provider = SimpleNamespace(model_name="stub", dimensions=4) + semantic = _semantic_search(index_name="milvus") matches = [ VectorMatch( key=VectorKey(entity_id=entity_id, chunk_key=f"entity:{entity_id}:0"), @@ -214,11 +219,11 @@ def hydrated_batch(_statement, params): session.execute.side_effect = hydrated_batch - hydrated = await repo._hydrate_vector_matches(session, matches) + hydrated = await semantic._hydrate_vector_matches(session, matches) assert session.execute.await_count == 3 assert max(len(call.args[1]) for call in session.execute.await_args_list) == 503 - assert [row["entity_id"] for row in hydrated] == list(range(600)) + assert [chunk.entity_id for chunk in hydrated] == list(range(600)) @pytest.mark.asyncio @@ -226,8 +231,6 @@ async def test_external_vector_query_overfetches_past_stale_adapter_hits( monkeypatch: pytest.MonkeyPatch, ) -> None: """Stale top-k extension hits must not crowd live manifest rows out.""" - repo = _ConcreteRepo() - repo._semantic_vector_index_name = "milvus" def matches(count: int) -> list[VectorMatch]: return [ @@ -239,15 +242,15 @@ def matches(count: int) -> list[VectorMatch]: ] adapter: Any = SimpleNamespace(search=AsyncMock(side_effect=[matches(2), matches(4)])) - repo._semantic_vector_index = adapter + semantic = _semantic_search(index_name="milvus", adapter=adapter) live_rows = [ - {"entity_id": 2, "chunk_key": "entity:2:0", "best_similarity": 0.9}, - {"entity_id": 3, "chunk_key": "entity:3:0", "best_similarity": 0.8}, + HydratedChunk(entity_id=2, chunk_key="entity:2:0", chunk_text="two", similarity=0.9), + HydratedChunk(entity_id=3, chunk_key="entity:3:0", chunk_text="three", similarity=0.8), ] hydrate = AsyncMock(side_effect=[[], live_rows]) - monkeypatch.setattr(repo, "_hydrate_vector_matches", hydrate) + monkeypatch.setattr(semantic, "_hydrate_vector_matches", hydrate) - result = await SearchRepositoryBase._run_vector_query(repo, AsyncMock(), [0.1], 2) + result = await semantic._run_vector_query(AsyncMock(), [0.1], 2) assert result == live_rows assert [call.kwargs["limit"] for call in adapter.search.await_args_list] == [2, 4] diff --git a/tests/repository/test_semantic_vector_sync.py b/tests/repository/test_semantic_vector_sync.py index fb1c12318..0e1003424 100644 --- a/tests/repository/test_semantic_vector_sync.py +++ b/tests/repository/test_semantic_vector_sync.py @@ -78,17 +78,6 @@ async def search( async def _ensure_vector_tables(self): pass - @override - async def _run_vector_query( - self, - session, - query_embedding, - candidate_limit, - *, - trace: SearchTraceCollector | None = None, - ): - return [] - @override async def _write_embeddings(self, session, jobs, embeddings): pass @@ -108,10 +97,6 @@ async def _delete_stale_chunks( ): return [] - @override - def _distance_to_similarity(self, distance: float) -> float: - return 1.0 / (1.0 + max(distance, 0.0)) - def _pending_job( entity_id: int = 1, diff --git a/tests/repository/test_vector_filter_candidate_restriction.py b/tests/repository/test_vector_filter_candidate_restriction.py index a40413862..8407a037a 100644 --- a/tests/repository/test_vector_filter_candidate_restriction.py +++ b/tests/repository/test_vector_filter_candidate_restriction.py @@ -30,10 +30,14 @@ from basic_memory import db from basic_memory.repository.embedding_provider import EmbeddingProvider from basic_memory.repository.search_filters import candidate_key_restriction_condition -from basic_memory.repository.search_repository_base import ( +from basic_memory.repository.search_query import PreparedSearchQuery +from basic_memory.repository.search_reader import ( VECTOR_FILTER_SCAN_LIMIT, VECTOR_HYDRATION_BATCH_SIZE, + HydratedChunk, + SemanticSearch, ) +from basic_memory.repository.semantic_vector_index import SemanticVectorIndex from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode # The scope the admitted rows share, plus a sibling scope the filter must reject. @@ -138,7 +142,11 @@ def _fake_embedding_provider() -> EmbeddingProvider: type( "EP", (), - {"embed_query": AsyncMock(return_value=[0.0] * 384), "dimensions": 384}, + { + "embed_query": AsyncMock(return_value=[0.0] * 384), + "dimensions": 384, + "model_name": "stub", + }, )(), ) @@ -148,18 +156,20 @@ def _semantic_repo(search_repository): search_repository._semantic_enabled = True search_repository._semantic_min_similarity = 0.0 search_repository._embedding_provider = _fake_embedding_provider() + # The nearest-neighbour stage is stubbed below, so the adapter is never consulted. + search_repository._semantic_vector_index = cast(SemanticVectorIndex, object()) return search_repository -def _vector_chunks(row_ids: list[int]) -> list[dict[str, Any]]: +def _vector_chunks(row_ids: list[int]) -> list[HydratedChunk]: """One vector hit per row, ranked in the order given.""" return [ - { - "chunk_key": f"{SearchItemType.ENTITY.value}:{row_id}:0", - "best_similarity": 0.99 - index * 0.001, - "chunk_text": TARGET_CONTENT, - "entity_id": row_id, - } + HydratedChunk( + entity_id=row_id, + chunk_key=f"{SearchItemType.ENTITY.value}:{row_id}:0", + chunk_text=TARGET_CONTENT, + similarity=0.99 - index * 0.001, + ) for index, row_id in enumerate(row_ids) ] @@ -174,9 +184,8 @@ async def test_filtered_vector_search_keeps_a_candidate_past_the_scan_window( with ( patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), - patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), patch.object( - repo, + SemanticSearch, "_run_vector_query", new_callable=AsyncMock, return_value=_vector_chunks([TARGET_ROW_ID]), @@ -210,36 +219,28 @@ async def test_filter_pass_answers_every_candidate_within_the_bind_bound( assert len(candidates) > VECTOR_HYDRATION_BATCH_SIZE batched_key_counts: list[int] = [] - original_search = repo.search + original_search = repo._fts.search - async def recording_search(*args, **kwargs): + async def recording_search(scope, query, **kwargs): if kwargs.get("candidate_keys") is not None: batched_key_counts.append(len(kwargs["candidate_keys"])) - return await original_search(*args, **kwargs) + return await original_search(scope, query, **kwargs) with ( - patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), - patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), patch.object( - repo, + SemanticSearch, "_run_vector_query", new_callable=AsyncMock, return_value=_vector_chunks(candidates), ), - patch.object(repo, "search", recording_search), + patch.object(repo._fts, "search", recording_search), ): - results = await repo._search_vector_only( - search_text="the answer", - permalink=None, - permalink_match=None, - title=None, - note_types=None, - after_date=None, - search_item_types=None, - categories=None, - metadata_filters=None, - file_path_prefix=SCOPE, - temporal=None, + results = await repo._semantic_search().vector_only( + PreparedSearchQuery( + search_text="the answer", + file_path_prefix=SCOPE, + retrieval_mode=SearchRetrievalMode.VECTOR, + ), limit=len(candidates), offset=0, ) diff --git a/tests/repository/test_vector_pagination.py b/tests/repository/test_vector_pagination.py index f2141ff05..ed79d7ce3 100644 --- a/tests/repository/test_vector_pagination.py +++ b/tests/repository/test_vector_pagination.py @@ -4,207 +4,43 @@ which requires a sufficiently large candidate_limit multiplier. """ -from sqlalchemy.ext.asyncio import AsyncSession -from basic_memory.repository.search_scope import ProjectScope -from collections.abc import Sequence -from contextlib import asynccontextmanager -from dataclasses import dataclass -from datetime import datetime -from typing import override, Any -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock import pytest -from basic_memory.repository.search_repository_base import ( - SearchIndexKey, - SearchRepositoryBase, -) -from basic_memory.repository.search_index_row import SearchIndexRow -from basic_memory.repository.search_trace import SearchTraceCollector -from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode -from basic_memory.temporal import TemporalFilter +from basic_memory.repository.search_reader import HydratedChunk +from tests.repository.test_vector_threshold import FakeRow, run_vector_only, vector_semantic -@dataclass -class FakeRow: - """Minimal stand-in for SearchIndexRow in pagination tests.""" - - id: int - type: str = "entity" - score: float = 0.0 - matched_chunk_text: str | None = None - content_snippet: str | None = None - - -class ConcreteSearchRepo(SearchRepositoryBase): - """Minimal concrete subclass for testing base class pagination logic.""" - - def __init__(self): - self._semantic_enabled = True - self._semantic_vector_k = 100 - self._semantic_min_similarity = 0.0 - self._embedding_provider = None - self._vector_dimensions = 384 - self._vector_tables_initialized = True - self.session_maker = None - self.project_id = 1 - self.scope = ProjectScope.single(1) - - @override - async def init_search_index(self): - pass # pragma: no cover - - @override - async def get_entity_physical_chunk_keys(self, entity_id: int) -> set[str] | None: - return None # physical storage is not inspectable in this double - - @override - async def search( - self, - search_text: str | None = None, - permalink: str | None = None, - permalink_match: str | None = None, - title: str | None = None, - note_types: list[str] | None = None, - after_date: datetime | None = None, - search_item_types: list[SearchItemType] | None = None, - categories: list[str] | None = None, - metadata_filters: dict[str, Any] | None = None, - file_path_prefix: str | None = None, - temporal: TemporalFilter | None = None, - retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, - min_similarity: float | None = None, - limit: int = 10, - offset: int = 0, - allow_relaxed: bool = False, - session: AsyncSession | None = None, - *, - candidate_keys: Sequence[SearchIndexKey] | None = None, - trace: SearchTraceCollector | None = None, - ) -> list[SearchIndexRow]: - return [] # pragma: no cover - - @override - async def _ensure_vector_tables(self): - pass # pragma: no cover - - @override - async def _run_vector_query( - self, - session, - query_embedding, - candidate_limit, - *, - trace: SearchTraceCollector | None = None, - ): - return [] # pragma: no cover - - @override - async def _write_embeddings(self, session, jobs, embeddings): - pass # pragma: no cover - - @override - async def _delete_entity_chunks(self, session, entity_id, *, expected_deletions=None): - return [] # pragma: no cover - - @override - async def _delete_stale_chunks( - self, - session, - stale_ids, - entity_id, - *, - expected_deletions=None, - ): - return [] # pragma: no cover - - async def _update_timestamp_sql(self): - return "CURRENT_TIMESTAMP" # pragma: no cover - - @override - def _distance_to_similarity(self, distance: float) -> float: - return 1.0 / (1.0 + max(distance, 0.0)) - - -@asynccontextmanager -async def fake_scoped_session(session_maker): - yield AsyncMock() - - -class _EmbeddingProvider: - dimensions = 384 - model_name = "stub" - - async def embed_query(self, text: str) -> list[float]: - return [0.0] * self.dimensions - - async def embed_documents(self, texts: list[str]) -> list[list[float]]: - return [[0.0] * self.dimensions for _ in texts] - - def runtime_log_attrs(self) -> dict[str, object]: - return {} - - -def _make_descending_vector_rows(count: int) -> list[dict[str, Any]]: - """Build vector rows with scores descending from ~1.0 to ~0.5.""" - rows = [] - for i in range(count): - # Similarity decreases linearly: 0.95, 0.94, 0.93, ... - similarity = 0.95 - (i * 0.01) - distance = (1.0 / similarity) - 1.0 - rows.append( - { - "chunk_key": f"entity:{i}:0", - "best_distance": distance, - "chunk_text": f"chunk text {i}", - } +def _make_descending_vector_rows(count: int) -> list[HydratedChunk]: + """Build vector rows with similarity descending from 0.95 in steps of 0.01.""" + return [ + HydratedChunk( + entity_id=index, + chunk_key=f"entity:{index}:0", + chunk_text=f"chunk text {index}", + similarity=0.95 - (index * 0.01), ) - return rows + for index in range(count) + ] @pytest.mark.asyncio async def test_page1_scores_gte_page2_scores(): """Page 1 minimum score must be >= page 2 maximum score.""" - repo = ConcreteSearchRepo() - + semantic = vector_semantic() # 20 results with descending scores fake_rows = _make_descending_vector_rows(20) - - repo._embedding_provider = _EmbeddingProvider() - fake_index_rows = {("entity", i): FakeRow(id=i) for i in range(20)} - async def run_page(offset, limit): - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", - fake_scoped_session, - ), - patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), - patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), - patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value=fake_index_rows, - ), - ): - return await repo._search_vector_only( - search_text="test", - permalink=None, - permalink_match=None, - title=None, - note_types=None, - after_date=None, - search_item_types=None, - categories=None, - metadata_filters=None, - file_path_prefix=None, - temporal=None, - limit=limit, - offset=offset, - ) + async def run_page(offset: int, limit: int): + return await run_vector_only( + semantic, + fake_rows, + AsyncMock(return_value=fake_index_rows), + limit=limit, + offset=offset, + ) page1 = await run_page(offset=0, limit=10) page2 = await run_page(offset=10, limit=10) diff --git a/tests/repository/test_vector_temporal_filter.py b/tests/repository/test_vector_temporal_filter.py index 0270d7815..0cf4a73cd 100644 --- a/tests/repository/test_vector_temporal_filter.py +++ b/tests/repository/test_vector_temporal_filter.py @@ -10,130 +10,88 @@ pin both halves at the seam rather than trusting the call sites to stay in step. """ -from typing import Any +from dataclasses import replace +from typing import Any, cast from unittest.mock import AsyncMock, patch import pytest +from basic_memory.repository.search_reader import SemanticSearch +from basic_memory.repository.search_scope import ProjectScope from basic_memory.temporal import TemporalFilter, TimeKind, parse_point from tests.repository.test_hybrid_fusion import ( - HYBRID_KWARGS, - ConcreteSearchRepo as HybridSearchRepo, + HYBRID_QUERY, + FakeFts, FakeRow as HybridFakeRow, + fake_vector_retrieval, ) from tests.repository.test_vector_threshold import ( - COMMON_SEARCH_KWARGS, - ConcreteSearchRepo as VectorSearchRepo, + VECTOR_QUERY, FakeRow, - _fake_embedding_provider, _make_vector_rows, - fake_scoped_session, + run_vector_only, + vector_semantic, ) TEMPORAL = TemporalFilter(kind=TimeKind.EFFECTIVE, at=parse_point("2026-07-28")) -def _vector_kwargs(**overrides: Any) -> dict[str, Any]: - return {**COMMON_SEARCH_KWARGS, **overrides} - - -def _hybrid_kwargs(**overrides: Any) -> dict[str, Any]: - return {**HYBRID_KWARGS, **overrides} - - -def _forwarded_temporal(leg: AsyncMock) -> Any: - """The `temporal` argument one retrieval leg was actually called with.""" - assert leg.await_args is not None, "leg was never awaited" - return leg.await_args.kwargs["temporal"] - - @pytest.mark.asyncio async def test_temporal_filter_applies_in_vector_mode(): """A valid-time filter narrows the vector candidate set, and is forwarded verbatim.""" - repo = VectorSearchRepo() - repo._semantic_min_similarity = 0.0 - repo._embedding_provider = _fake_embedding_provider(AsyncMock(return_value=[0.0] * 384)) - # The embedding neighbourhood offers three entities; only entity 1 asserts a range # covering the queried date, so the FTS intersection pass returns just that one. - filter_pass = AsyncMock(return_value=[FakeRow(id=1)]) - - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), - patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), - patch.object( - repo, - "_run_vector_query", - new_callable=AsyncMock, - return_value=_make_vector_rows([0.9, 0.8, 0.7]), - ), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value={("entity", i): FakeRow(id=i) for i in range(3)}, - ), - patch.object(repo, "search", filter_pass), - ): - results = await repo._search_vector_only(**_vector_kwargs(temporal=TEMPORAL)) + filter_pass = FakeFts([FakeRow(id=1)]) + semantic = vector_semantic(fts=filter_pass) + + results = await run_vector_only( + semantic, + _make_vector_rows([0.9, 0.8, 0.7]), + AsyncMock(return_value={("entity", i): FakeRow(id=i) for i in range(3)}), + query=replace(VECTOR_QUERY, temporal=TEMPORAL), + ) assert [row.id for row in results] == [1] # Counted as a requested filter... - filter_pass.assert_awaited_once() - # ...and forwarded unchanged, so the intersection asks the same question. - assert _forwarded_temporal(filter_pass) is TEMPORAL + assert len(filter_pass.queries) == 1 + # ...and forwarded unchanged, so the intersection asks the same question, + # about the candidates themselves rather than a page of the whole match set. + assert filter_pass.queries[0].temporal is TEMPORAL + assert filter_pass.queries[0].search_text is None + assert filter_pass.calls[0]["candidate_keys"] == [("entity", 0), ("entity", 1), ("entity", 2)] @pytest.mark.asyncio async def test_vector_mode_without_a_temporal_filter_runs_no_intersection_pass(): """An unfiltered semantic search must not pay for a filter pass it does not need.""" - repo = VectorSearchRepo() - repo._semantic_min_similarity = 0.0 - repo._embedding_provider = _fake_embedding_provider(AsyncMock(return_value=[0.0] * 384)) - filter_pass = AsyncMock(return_value=[]) - - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), - patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), - patch.object( - repo, - "_run_vector_query", - new_callable=AsyncMock, - return_value=_make_vector_rows([0.9]), - ), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value={("entity", 0): FakeRow(id=0)}, - ), - patch.object(repo, "search", filter_pass), - ): - results = await repo._search_vector_only(**_vector_kwargs()) + filter_pass = FakeFts([]) + semantic = vector_semantic(fts=filter_pass) + + results = await run_vector_only( + semantic, + _make_vector_rows([0.9]), + AsyncMock(return_value={("entity", 0): FakeRow(id=0)}), + ) assert [row.id for row in results] == [0] - filter_pass.assert_not_awaited() + assert filter_pass.queries == [] @pytest.mark.asyncio async def test_temporal_filter_applies_in_hybrid_mode(): """Hybrid fuses two legs; both must ask the same valid-time question.""" - repo = HybridSearchRepo() - fts_leg = AsyncMock(return_value=[HybridFakeRow(id=1, score=5.0, title="dated")]) + fts_leg = FakeFts([HybridFakeRow(id=1, score=5.0, title="dated")]) + semantic = SemanticSearch( + cast(Any, None), ProjectScope.single(1), fts_leg, fake_vector_retrieval() + ) vector_leg = AsyncMock(return_value=[HybridFakeRow(id=1, score=0.9, title="dated")]) - with ( - patch.object(repo, "search", fts_leg), - patch.object(repo, "_search_vector_only", vector_leg), - ): - results = await repo._search_hybrid(**_hybrid_kwargs(temporal=TEMPORAL)) + with patch.object(semantic, "vector_only", vector_leg): + results = await semantic.hybrid( + replace(HYBRID_QUERY, temporal=TEMPORAL), limit=10, offset=0 + ) assert [row.id for row in results] == [1] - assert _forwarded_temporal(fts_leg) is TEMPORAL - assert _forwarded_temporal(vector_leg) is TEMPORAL + assert fts_leg.queries[0].temporal is TEMPORAL + assert vector_leg.await_args is not None, "vector leg was never awaited" + assert vector_leg.await_args.args[0].temporal is TEMPORAL diff --git a/tests/repository/test_vector_threshold.py b/tests/repository/test_vector_threshold.py index fdfb6d8a7..4314176d5 100644 --- a/tests/repository/test_vector_threshold.py +++ b/tests/repository/test_vector_threshold.py @@ -1,27 +1,22 @@ """Tests for semantic_min_similarity threshold filtering in vector search.""" -from sqlalchemy.ext.asyncio import AsyncSession -from basic_memory.repository.search_scope import ProjectScope -from collections.abc import Sequence from contextlib import asynccontextmanager -from dataclasses import dataclass -from datetime import datetime -from typing import override, Any, Optional, cast +from dataclasses import dataclass, replace +from typing import Any from unittest.mock import AsyncMock, patch import pytest -from basic_memory.repository.embedding_provider import EmbeddingProvider -from basic_memory.repository.search_index_row import SearchIndexRow -from basic_memory.repository.search_repository_base import ( +from basic_memory.repository.search_query import PreparedSearchQuery +from basic_memory.repository.search_reader import ( SMALL_NOTE_CONTENT_LIMIT, TOP_CHUNKS_PER_RESULT, - SearchIndexKey, - SearchRepositoryBase, + HydratedChunk, + SemanticSearch, ) -from basic_memory.repository.search_trace import SearchTraceCollector -from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode -from basic_memory.temporal import TemporalFilter +from basic_memory.repository.search_scope import ProjectScope +from basic_memory.schemas.search import SearchRetrievalMode +from tests.repository.test_hybrid_fusion import FakeFts, fake_vector_retrieval @dataclass @@ -35,174 +30,71 @@ class FakeRow: content_snippet: str | None = None -class ConcreteSearchRepo(SearchRepositoryBase): - """Minimal concrete subclass for testing base class threshold logic.""" - - def __init__(self): - # Skip super().__init__ — we only need the attributes under test - self._semantic_enabled = True - self._semantic_vector_k = 100 - self._semantic_min_similarity = 0.0 - self._embedding_provider = None - self._vector_dimensions = 384 - self._vector_tables_initialized = True - self.session_maker = None - self.project_id = 1 - self.scope = ProjectScope.single(1) - - # --- Abstract method stubs (not exercised by these tests) --- - - @override - async def init_search_index(self): - pass # pragma: no cover - - @override - async def get_entity_physical_chunk_keys(self, entity_id: int) -> set[str] | None: - return None # physical storage is not inspectable in this double - - @override - async def search( - self, - search_text: Optional[str] = None, - permalink: Optional[str] = None, - permalink_match: Optional[str] = None, - title: Optional[str] = None, - note_types: Optional[list[str]] = None, - after_date: Optional[datetime] = None, - search_item_types: Optional[list[SearchItemType]] = None, - categories: Optional[list[str]] = None, - metadata_filters: Optional[dict[str, Any]] = None, - file_path_prefix: Optional[str] = None, - temporal: Optional[TemporalFilter] = None, - retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, - min_similarity: Optional[float] = None, - limit: int = 10, - offset: int = 0, - allow_relaxed: bool = False, - session: AsyncSession | None = None, - *, - candidate_keys: Sequence[SearchIndexKey] | None = None, - trace: SearchTraceCollector | None = None, - ) -> list[SearchIndexRow]: - return [] # pragma: no cover - - @override - async def _ensure_vector_tables(self): - pass # pragma: no cover - - @override - async def _run_vector_query( - self, - session, - query_embedding, - candidate_limit, - *, - trace: SearchTraceCollector | None = None, - ): - return [] # pragma: no cover - - @override - async def _write_embeddings(self, session, jobs, embeddings): - pass # pragma: no cover - - @override - async def _delete_entity_chunks(self, session, entity_id, *, expected_deletions=None): - return [] # pragma: no cover - - @override - async def _delete_stale_chunks( - self, - session, - stale_ids, - entity_id, - *, - expected_deletions=None, - ): - return [] # pragma: no cover - - async def _update_timestamp_sql(self): - return "CURRENT_TIMESTAMP" # pragma: no cover - - @override - def _distance_to_similarity(self, distance: float) -> float: - return 1.0 / (1.0 + max(distance, 0.0)) - - -def _make_vector_rows(scores: list[float]) -> list[dict[str, Any]]: - """Build fake vector query rows with controlled distances. - - Distance = (1/score) - 1 inverts the similarity formula: - similarity = 1 / (1 + distance) - """ - rows = [] - for i, score in enumerate(scores): - distance = (1.0 / score) - 1.0 - rows.append( - { - "chunk_key": f"entity:{i}:0", - "best_distance": distance, - "chunk_text": f"chunk text for entity:{i}:0", - } +def _make_vector_rows(scores: list[float]) -> list[HydratedChunk]: + """One hydrated chunk per search row, ranked at the given similarity.""" + return [ + HydratedChunk( + entity_id=index, + chunk_key=f"entity:{index}:0", + chunk_text=f"chunk text for entity:{index}:0", + similarity=score, ) - return rows + for index, score in enumerate(scores) + ] -def _fake_embedding_provider(mock_embed: AsyncMock) -> EmbeddingProvider: - return cast( - EmbeddingProvider, - type("EP", (), {"embed_query": mock_embed, "dimensions": 384})(), - ) +def fake_session_maker() -> Any: + """A session factory for the hydration step; the stubbed stages never touch it.""" + @asynccontextmanager + async def session(): + yield AsyncMock() -@asynccontextmanager -async def fake_scoped_session(session_maker): - """Fake scoped_session that yields a mock session object.""" - yield AsyncMock() - - -COMMON_SEARCH_KWARGS: dict[str, Any] = dict( - search_text="test", - permalink=None, - permalink_match=None, - title=None, - note_types=None, - after_date=None, - search_item_types=None, - categories=None, - metadata_filters=None, - file_path_prefix=None, - temporal=None, - limit=10, - offset=0, -) + return session -@pytest.mark.asyncio -async def test_threshold_zero_returns_all(): - """With threshold=0.0 (default), all results pass through.""" - repo = ConcreteSearchRepo() - repo._semantic_min_similarity = 0.0 +def vector_semantic(*, min_similarity: float = 0.0, fts: FakeFts | None = None) -> SemanticSearch: + """A vector pipeline with a stubbed adapter, ready for its neighbour stage to be patched.""" + return SemanticSearch( + fake_session_maker(), + ProjectScope.single(1), + fts or FakeFts(), + fake_vector_retrieval(min_similarity=min_similarity), + ) + - fake_rows = _make_vector_rows([0.9, 0.5, 0.3]) +VECTOR_QUERY = PreparedSearchQuery(search_text="test", retrieval_mode=SearchRetrievalMode.VECTOR) - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) +async def run_vector_only( + semantic: SemanticSearch, + vector_rows: list[HydratedChunk], + fetch_rows: AsyncMock, + *, + query: PreparedSearchQuery = VECTOR_QUERY, + limit: int = 10, + offset: int = 0, +) -> list[Any]: + """Run vector-only search with the neighbour and row-fetch stages stubbed.""" with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), - patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), - patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows), patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value={("entity", i): FakeRow(id=i) for i in range(3)}, + semantic, "_run_vector_query", new_callable=AsyncMock, return_value=vector_rows ), + patch.object(semantic, "_fetch_search_index_rows_by_ids", fetch_rows), ): - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS) + return await semantic.vector_only(query, limit=limit, offset=offset) + + +def _index_rows(count: int) -> AsyncMock: + return AsyncMock(return_value={("entity", i): FakeRow(id=i) for i in range(count)}) + + +@pytest.mark.asyncio +async def test_threshold_zero_returns_all(): + """With threshold=0.0 (default), all results pass through.""" + semantic = vector_semantic(min_similarity=0.0) + + results = await run_vector_only(semantic, _make_vector_rows([0.9, 0.5, 0.3]), _index_rows(3)) assert len(results) == 3 @@ -210,129 +102,56 @@ async def test_threshold_zero_returns_all(): @pytest.mark.asyncio async def test_threshold_filters_low_scores(): """Results below the threshold are excluded.""" - repo = ConcreteSearchRepo() - repo._semantic_min_similarity = 0.6 - - # Scores: 0.9 (pass), 0.5 (fail), 0.3 (fail) - fake_rows = _make_vector_rows([0.9, 0.5, 0.3]) - - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) + semantic = vector_semantic(min_similarity=0.6) - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), - patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), - patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - # Only entity_0 (score=0.9) passes the threshold; the fetch only gets id 0 - return_value={("entity", 0): FakeRow(id=0)}, - ), - ): - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS) + # Scores: 0.9 (pass), 0.5 (fail), 0.3 (fail). Only entity_0 reaches the row fetch. + results = await run_vector_only(semantic, _make_vector_rows([0.9, 0.5, 0.3]), _index_rows(1)) - # Only the 0.9 result passes the 0.6 threshold assert len(results) == 1 @pytest.mark.asyncio async def test_threshold_returns_empty_when_all_below(): """All results below threshold → empty list, no DB fetch.""" - repo = ConcreteSearchRepo() - repo._semantic_min_similarity = 0.8 - - # All scores below 0.8 - fake_rows = _make_vector_rows([0.5, 0.4, 0.3]) + semantic = vector_semantic(min_similarity=0.8) + fetch_rows = AsyncMock() - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) - - mock_fetch = AsyncMock() - - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), - patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), - patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows), - patch.object(repo, "_fetch_search_index_rows_by_ids", mock_fetch), - ): - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS) + results = await run_vector_only(semantic, _make_vector_rows([0.5, 0.4, 0.3]), fetch_rows) assert results == [] # Should short-circuit before fetching search_index rows - mock_fetch.assert_not_called() + fetch_rows.assert_not_called() @pytest.mark.asyncio -async def test_per_query_min_similarity_overrides_instance_default(): - """Per-query min_similarity takes precedence over instance-level default.""" - repo = ConcreteSearchRepo() - # Instance default would filter out 0.5 and 0.3 - repo._semantic_min_similarity = 0.6 - - # Scores: 0.9, 0.5, 0.3 - fake_rows = _make_vector_rows([0.9, 0.5, 0.3]) - - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) - - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), - patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), - patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value={("entity", i): FakeRow(id=i) for i in range(3)}, - ), - ): - # Override to 0.0 → all results pass through despite instance default of 0.6 - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS, min_similarity=0.0) +async def test_per_query_min_similarity_overrides_configured_default(): + """Per-query min_similarity takes precedence over the configured default.""" + # The configured default would filter out 0.5 and 0.3 + semantic = vector_semantic(min_similarity=0.6) + + # Override to 0.0 → all results pass through despite the configured 0.6 + results = await run_vector_only( + semantic, + _make_vector_rows([0.9, 0.5, 0.3]), + _index_rows(3), + query=replace(VECTOR_QUERY, min_similarity=0.0), + ) assert len(results) == 3 @pytest.mark.asyncio async def test_per_query_min_similarity_tightens_threshold(): - """Per-query min_similarity=0.8 filters more aggressively than instance default.""" - repo = ConcreteSearchRepo() - # Instance default is permissive - repo._semantic_min_similarity = 0.0 - - # Scores: 0.9, 0.5, 0.3 - fake_rows = _make_vector_rows([0.9, 0.5, 0.3]) - - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) - - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), - patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), - patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - # Only id=0 (score=0.9) will be fetched after filtering - return_value={("entity", 0): FakeRow(id=0)}, - ), - ): - # Override to 0.8 → only score=0.9 passes - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS, min_similarity=0.8) + """Per-query min_similarity=0.8 filters more aggressively than the configured default.""" + semantic = vector_semantic(min_similarity=0.0) + + # Override to 0.8 → only score=0.9 passes + results = await run_vector_only( + semantic, + _make_vector_rows([0.9, 0.5, 0.3]), + _index_rows(1), + query=replace(VECTOR_QUERY, min_similarity=0.8), + ) assert len(results) == 1 assert results[0].id == 0 @@ -341,29 +160,9 @@ async def test_per_query_min_similarity_tightens_threshold(): @pytest.mark.asyncio async def test_matched_chunk_text_populated_on_vector_results(): """Vector search results carry the matched chunk text from the best-matching chunk.""" - repo = ConcreteSearchRepo() - repo._semantic_min_similarity = 0.0 - - fake_rows = _make_vector_rows([0.9, 0.7]) - - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) + semantic = vector_semantic() - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), - patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), - patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value={("entity", i): FakeRow(id=i) for i in range(2)}, - ), - ): - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS) + results = await run_vector_only(semantic, _make_vector_rows([0.9, 0.7]), _index_rows(2)) assert len(results) == 2 # Results are sorted by score descending, so id=0 (0.9) first, id=1 (0.7) second @@ -372,56 +171,32 @@ async def test_matched_chunk_text_populated_on_vector_results(): assert results[1].matched_chunk_text == "chunk text for entity:1:0" -def _make_multi_chunk_vector_rows(si_id: int, scores: list[float]) -> list[dict[str, Any]]: - """Build multiple fake vector chunks for a single search_index row. - - Each chunk gets a unique chunk_index within the same si_id. - Distance = (1/score) - 1 inverts the similarity formula. - """ - rows = [] - for chunk_idx, score in enumerate(scores): - distance = (1.0 / score) - 1.0 - rows.append( - { - "chunk_key": f"entity:{si_id}:{chunk_idx}", - "best_distance": distance, - "chunk_text": f"chunk-{chunk_idx} (sim={score})", - } +def _make_multi_chunk_vector_rows(si_id: int, scores: list[float]) -> list[HydratedChunk]: + """Several chunks of one search row, each at its own similarity.""" + return [ + HydratedChunk( + entity_id=si_id, + chunk_key=f"entity:{si_id}:{chunk_index}", + chunk_text=f"chunk-{chunk_index} (sim={score})", + similarity=score, ) - return rows + for chunk_index, score in enumerate(scores) + ] @pytest.mark.asyncio async def test_top_n_chunks_joined_in_matched_chunk_text(): """Large note with 7 chunks: top 5 by similarity are joined with separator.""" - repo = ConcreteSearchRepo() - repo._semantic_min_similarity = 0.0 - - # 7 chunks for entity 0, with varying similarities + semantic = vector_semantic() chunk_scores = [0.6, 0.9, 0.4, 0.8, 0.75, 0.3, 0.85] - fake_rows = _make_multi_chunk_vector_rows(si_id=0, scores=chunk_scores) - - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) - # content_snippet exceeds SMALL_NOTE_CONTENT_LIMIT → top-N chunks path large_content = "x" * (SMALL_NOTE_CONTENT_LIMIT + 1) - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), - patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), - patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value={("entity", 0): FakeRow(id=0, content_snippet=large_content)}, - ), - ): - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS) + results = await run_vector_only( + semantic, + _make_multi_chunk_vector_rows(si_id=0, scores=chunk_scores), + AsyncMock(return_value={("entity", 0): FakeRow(id=0, content_snippet=large_content)}), + ) assert len(results) == 1 text = results[0].matched_chunk_text @@ -440,32 +215,15 @@ async def test_top_n_chunks_joined_in_matched_chunk_text(): @pytest.mark.asyncio async def test_small_note_returns_full_content_as_matched_chunk(): """Small note (content_snippet under limit) returns full content instead of chunks.""" - repo = ConcreteSearchRepo() - repo._semantic_min_similarity = 0.0 - - fake_rows = _make_vector_rows([0.9]) - - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) - + semantic = vector_semantic() small_content = "This is a short note with all the important details." assert len(small_content) <= SMALL_NOTE_CONTENT_LIMIT - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), - patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), - patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value={("entity", 0): FakeRow(id=0, content_snippet=small_content)}, - ), - ): - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS) + results = await run_vector_only( + semantic, + _make_vector_rows([0.9]), + AsyncMock(return_value={("entity", 0): FakeRow(id=0, content_snippet=small_content)}), + ) assert len(results) == 1 # Full content returned instead of the chunk text @@ -475,33 +233,31 @@ async def test_small_note_returns_full_content_as_matched_chunk(): @pytest.mark.asyncio async def test_large_note_returns_chunks_not_full_content(): """Large note (content_snippet over limit) returns top-N chunks, not full content.""" - repo = ConcreteSearchRepo() - repo._semantic_min_similarity = 0.0 - - fake_rows = _make_vector_rows([0.9]) - - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) - + semantic = vector_semantic() large_content = "x" * (SMALL_NOTE_CONTENT_LIMIT + 500) - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), - patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), - patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value={("entity", 0): FakeRow(id=0, content_snippet=large_content)}, - ), - ): - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS) + results = await run_vector_only( + semantic, + _make_vector_rows([0.9]), + AsyncMock(return_value={("entity", 0): FakeRow(id=0, content_snippet=large_content)}), + ) assert len(results) == 1 # Should use chunk text, not the full content assert results[0].matched_chunk_text == "chunk text for entity:0:0" assert results[0].matched_chunk_text != large_content + + +@pytest.mark.asyncio +async def test_unparseable_chunk_key_names_no_search_row(): + """A chunk whose key does not spell a search row is skipped rather than ranked.""" + semantic = vector_semantic() + rows = [ + HydratedChunk(entity_id=0, chunk_key="entity:0:0", chunk_text="good", similarity=0.9), + HydratedChunk(entity_id=0, chunk_key="garbage", chunk_text="bad", similarity=0.95), + ] + + results = await run_vector_only(semantic, rows, _index_rows(1)) + + assert [row.id for row in results] == [0] + assert results[0].matched_chunk_text == "good" diff --git a/tests/services/test_project_readiness.py b/tests/services/test_project_readiness.py index a7aebf5d7..277831cd9 100644 --- a/tests/services/test_project_readiness.py +++ b/tests/services/test_project_readiness.py @@ -16,7 +16,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from basic_memory.repository.search_repository import create_search_repository -from basic_memory.repository.search_repository_base import VECTOR_HYDRATION_BATCH_SIZE +from basic_memory.repository.search_reader import VECTOR_HYDRATION_BATCH_SIZE from basic_memory.repository.embedding_provider_factory import ( configured_embedding_provider_identity, ) From 2386890bffd7b664ed17737e3be44a723632e501 Mon Sep 17 00:00:00 2001 From: Paul Hernandez <60959+phernandez@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:55:34 -0500 Subject: [PATCH 4/8] refactor(core): bind vector adapters to the database and search them over a ProjectScope (#1566) Signed-off-by: phernandez Co-authored-by: Claude Fable 5.1 --- CHANGELOG.md | 9 + src/basic_memory/repository/milvus_index.py | 128 ++++++------ src/basic_memory/repository/pgvector_index.py | 75 +++---- .../repository/postgres_search_repository.py | 12 +- src/basic_memory/repository/search_reader.py | 8 +- .../repository/search_repository.py | 1 - .../repository/search_repository_base.py | 14 +- .../repository/semantic_vector_index.py | 40 ++-- .../semantic_vector_index_factory.py | 10 +- .../repository/sqlite_search_repository.py | 12 +- .../repository/sqlite_vec_index.py | 68 ++---- test-int/semantic/conftest.py | 1 - test-int/semantic/test_milvus_lite.py | 33 +-- test-int/test_embedding_status_vec0.py | 6 +- tests/repository/test_milvus_index.py | 194 ++++++++++++------ tests/repository/test_pgvector_index.py | 55 +++-- tests/repository/test_search_reader.py | 2 +- tests/repository/test_search_trace.py | 15 +- tests/repository/test_semantic_search_base.py | 28 ++- .../repository/test_semantic_vector_index.py | 25 +-- .../test_sqlite_vector_search_repository.py | 88 +++++++- ...st_vector_manifest_generation_ownership.py | 9 +- 22 files changed, 484 insertions(+), 349 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8b57ec9a5..0d347cd53 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -143,6 +143,15 @@ construct the pipeline directly instead of subclassing the repository. No query behavior changes. +- **#1558**: Vector adapters are bound to the database, not to a project. `VectorIndexScope` + is the database namespace plus embedding schema; every write names the project it + touches (`upsert(project_id, ...)`, `delete(project_id, ...)`, + `delete_entity(project_id, ...)`, `delete_orphans(project_id, ...)`) and `search` takes + a `ProjectScope`, so one sqlite-vec or pgvector adapter answers a query across any set + of projects with one statement. Milvus keeps a collection per project and searches the + collections in scope. `SemanticSearch` passes its scope through, so a project + repository's vector search is unchanged. No query behavior changes. + ## v0.23.2 (2026-08-25) diff --git a/src/basic_memory/repository/milvus_index.py b/src/basic_memory/repository/milvus_index.py index 9cc367c15..dcc8e6ce2 100644 --- a/src/basic_memory/repository/milvus_index.py +++ b/src/basic_memory/repository/milvus_index.py @@ -13,6 +13,7 @@ MilvusStoredRecord, create_repository, ) +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_vector_index import ( VectorDeletion, VectorIndexScope, @@ -33,10 +34,10 @@ def _record_id(key: VectorKey) -> str: return hashlib.sha256(stable_key).hexdigest() -def collection_name(settings: MilvusSettings, scope: VectorIndexScope) -> str: - """Return a stable project collection name independent of embedding schema.""" +def collection_name(settings: MilvusSettings, scope: VectorIndexScope, project_id: int) -> str: + """Return one project's stable collection name, independent of embedding schema.""" namespace_digest = hashlib.sha256(scope.namespace.encode()).hexdigest()[:24] - return f"{settings.collection_prefix}_{namespace_digest}_{scope.project_id}" + return f"{settings.collection_prefix}_{namespace_digest}_{project_id}" def _normalize_cosine_score(score: float) -> float: @@ -45,7 +46,7 @@ def _normalize_cosine_score(score: float) -> float: class MilvusVectorIndex: - """Persist and query one Basic Memory project's vectors in Milvus.""" + """Persist and query Basic Memory vectors in Milvus, one collection per project.""" def __init__( self, @@ -56,10 +57,10 @@ def __init__( ) -> None: self.scope = scope self._settings = settings - self._collection_name = collection_name(settings, scope) self._repository_factory = repository_factory - self._initialized = False - self._initialize_lock = asyncio.Lock() + # Projects whose collection has been created or validated by this instance. + self._ready_projects: set[int] = set() + self._collection_lock = asyncio.Lock() def _with_repository[T](self, operation: Callable[[MilvusRepository], T]) -> T: repository = self._repository_factory(self._settings) @@ -68,27 +69,27 @@ def _with_repository[T](self, operation: Callable[[MilvusRepository], T]) -> T: finally: repository.close() - def _initialize_blocking(self) -> None: + def _collection(self, project_id: int) -> str: + return collection_name(self._settings, self.scope, project_id) + + def _validate_collection_blocking(self, collection: str) -> None: def initialize_repository(repository: MilvusRepository) -> None: - dimensions = repository.collection_dimensions(self._collection_name) + dimensions = repository.collection_dimensions(collection) if dimensions is None: - created = repository.create_collection( - self._collection_name, - self.scope.dimensions, - ) + created = repository.create_collection(collection, self.scope.dimensions) if created: return - dimensions = repository.collection_dimensions(self._collection_name) + dimensions = repository.collection_dimensions(collection) if dimensions is None: raise RuntimeError( - f"Milvus collection '{self._collection_name}' disappeared after " + f"Milvus collection '{collection}' disappeared after " "a concurrent create operation." ) if dimensions == self.scope.dimensions: # Milvus Lite releases persisted collections when the owning process exits. # Load only after the scope check so migrations do not load incompatible # remote collections before Basic Memory refuses to use them. - repository.load_collection(self._collection_name) + repository.load_collection(collection) return # Trigger: an existing project collection uses another embedding dimension. @@ -96,7 +97,7 @@ def initialize_repository(repository: MilvusRepository) -> None: # repeatedly erase each other's vectors during a rolling deployment. # Outcome: preserve the collection until an operator coordinates migration. raise RuntimeError( - f"Milvus collection '{self._collection_name}' has {dimensions} dimensions, " + f"Milvus collection '{collection}' has {dimensions} dimensions, " f"but Basic Memory is configured for {self.scope.dimensions}. Refusing to " "replace shared vector storage automatically; stop all writers and coordinate " "the collection migration before reindexing." @@ -106,12 +107,13 @@ def initialize_repository(repository: MilvusRepository) -> None: def _search_blocking( self, + collection: str, query: Sequence[float], limit: int, ) -> list[MilvusStoredMatch]: repository = self._repository_factory(self._settings) try: - return repository.search(self._collection_name, query, limit) + return repository.search(collection, query, limit) finally: repository.close() @@ -135,19 +137,28 @@ async def _run_blocking_mutation(self, operation: Callable[[], None]) -> None: raise async def initialize(self) -> None: - if self._initialized: - return - async with self._initialize_lock: - if self._initialized: - return - await self._run_blocking_mutation(self._initialize_blocking) - self._initialized = True + """Nothing is shared across projects: each collection is validated on first use.""" + return None + + async def _ensure_collection(self, project_id: int) -> str: + """Create or validate one project's collection once per adapter instance.""" + collection = self._collection(project_id) + if project_id in self._ready_projects: + return collection + async with self._collection_lock: + if project_id in self._ready_projects: + return collection + await self._run_blocking_mutation( + lambda: self._validate_collection_blocking(collection) + ) + self._ready_projects.add(project_id) + return collection - async def upsert(self, records: Sequence[VectorRecord]) -> None: + async def upsert(self, project_id: int, records: Sequence[VectorRecord]) -> None: if not records: return validate_vector_dimensions(self.scope, records) - await self.initialize() + collection = await self._ensure_collection(project_id) stored_records = [ MilvusStoredRecord( record_id=_record_id(record.key), @@ -160,47 +171,44 @@ async def upsert(self, records: Sequence[VectorRecord]) -> None: ] await self._run_blocking_mutation( lambda: self._with_repository( - lambda repository: repository.upsert(self._collection_name, stored_records) + lambda repository: repository.upsert(collection, stored_records) ) ) - async def delete(self, records: Sequence[VectorDeletion]) -> None: + async def delete(self, project_id: int, records: Sequence[VectorDeletion]) -> None: if not records: return - await self.initialize() + collection = await self._ensure_collection(project_id) stored_deletions = [(_record_id(record.key), record.source_hash) for record in records] await self._run_blocking_mutation( lambda: self._with_repository( - lambda repository: repository.delete_records( - self._collection_name, - stored_deletions, - ) + lambda repository: repository.delete_records(collection, stored_deletions) ) ) - async def delete_entity(self, entity_id: int) -> None: - await self.initialize() + async def delete_entity(self, project_id: int, entity_id: int) -> None: + collection = await self._ensure_collection(project_id) await self._run_blocking_mutation( lambda: self._with_repository( - lambda repository: repository.delete_entity(self._collection_name, entity_id) + lambda repository: repository.delete_entity(collection, entity_id) ) ) - async def delete_orphans(self, live_keys: Sequence[VectorKey]) -> None: - await self.initialize() + async def delete_orphans(self, project_id: int, live_keys: Sequence[VectorKey]) -> None: + collection = await self._ensure_collection(project_id) live_ids = {_record_id(key) for key in live_keys} def delete_missing(repository: MilvusRepository) -> None: orphan_ids: list[str] = [] - for record_id in repository.iter_ids(self._collection_name): + for record_id in repository.iter_ids(collection): if record_id in live_ids: continue orphan_ids.append(record_id) if len(orphan_ids) == _ORPHAN_DELETE_BATCH_SIZE: - repository.delete_ids(self._collection_name, orphan_ids) + repository.delete_ids(collection, orphan_ids) orphan_ids.clear() if orphan_ids: - repository.delete_ids(self._collection_name, orphan_ids) + repository.delete_ids(collection, orphan_ids) await self._run_blocking_mutation(lambda: self._with_repository(delete_missing)) @@ -209,28 +217,32 @@ async def search( query: Sequence[float], *, limit: int, + projects: ProjectScope, ) -> list[VectorMatch]: - if not query or limit <= 0: + if not query or limit <= 0 or projects.is_empty: return [] validate_query_dimensions(self.scope, query) - await self.initialize() - - stored_matches = await asyncio.to_thread(self._search_blocking, query, limit) - matches = [ - VectorMatch( - key=VectorKey( - entity_id=match.entity_id, - chunk_key=match.chunk_key, - ), - similarity=_normalize_cosine_score(match.score), + + # Milvus has no cross-collection search, so a scope wider than one project + # asks each project's collection for its own top ``limit`` and merges them. + matches: list[VectorMatch] = [] + for project_id in projects.project_ids: + collection = await self._ensure_collection(project_id) + stored_matches = await asyncio.to_thread( + self._search_blocking, collection, query, limit ) - for match in stored_matches - ] - return sorted( - matches, + matches.extend( + VectorMatch( + key=VectorKey(entity_id=match.entity_id, chunk_key=match.chunk_key), + similarity=_normalize_cosine_score(match.score), + ) + for match in stored_matches + ) + matches.sort( key=lambda match: ( -match.similarity, match.key.entity_id, match.key.chunk_key, - ), + ) ) + return matches[:limit] diff --git a/src/basic_memory/repository/pgvector_index.py b/src/basic_memory/repository/pgvector_index.py index 764a1c716..1db0be8c5 100644 --- a/src/basic_memory/repository/pgvector_index.py +++ b/src/basic_memory/repository/pgvector_index.py @@ -10,6 +10,7 @@ from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from basic_memory import db +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError from basic_memory.repository.semantic_vector_index import ( VectorDeletion, @@ -151,37 +152,7 @@ async def _has_source_hash_column(self, session: AsyncSession) -> bool: ) return result.scalar_one_or_none() is not None - async def _chunk_ids_by_key( - self, - session: AsyncSession, - keys: Sequence[VectorKey], - ) -> dict[VectorKey, int]: - if not keys: - return {} - - params: dict[str, object] = {"project_id": self.scope.project_id} - predicates: list[str] = [] - for index, key in enumerate(keys): - params[f"entity_id_{index}"] = key.entity_id - params[f"chunk_key_{index}"] = key.chunk_key - predicates.append( - f"(entity_id = :entity_id_{index} AND chunk_key = :chunk_key_{index})" - ) - result = await session.execute( - text( - "SELECT id, entity_id, chunk_key FROM search_vector_chunks " - "WHERE project_id = :project_id AND (" + " OR ".join(predicates) + ")" - ), - params, - ) - return { - VectorKey(entity_id=int(row["entity_id"]), chunk_key=str(row["chunk_key"])): int( - row["id"] - ) - for row in result.mappings().all() - } - - async def upsert(self, records: Sequence[VectorRecord]) -> None: + async def upsert(self, project_id: int, records: Sequence[VectorRecord]) -> None: if not records: return validate_vector_dimensions(self.scope, records) @@ -189,7 +160,7 @@ async def upsert(self, records: Sequence[VectorRecord]) -> None: async with db.scoped_session(self._session_maker) as session: keys = [record.key for record in records] - params: dict[str, object] = {"project_id": self.scope.project_id} + params: dict[str, object] = {"project_id": project_id} predicates: list[str] = [] for index, key in enumerate(keys): params[f"entity_id_{index}"] = key.entity_id @@ -224,7 +195,7 @@ async def upsert(self, records: Sequence[VectorRecord]) -> None: if not current_records: return - params = {"project_id": self.scope.project_id} + params = {"project_id": project_id} values: list[str] = [] for index, record in enumerate(current_records): params[f"chunk_id_{index}"] = manifest_by_key[record.key][0] @@ -252,12 +223,12 @@ async def upsert(self, records: Sequence[VectorRecord]) -> None: ) await session.commit() - async def delete(self, records: Sequence[VectorDeletion]) -> None: + async def delete(self, project_id: int, records: Sequence[VectorDeletion]) -> None: if not records: return await self.initialize() async with db.scoped_session(self._session_maker) as session: - params: dict[str, object] = {"project_id": self.scope.project_id} + params: dict[str, object] = {"project_id": project_id} predicates: list[str] = [] for index, record in enumerate(records): params[f"entity_id_{index}"] = record.key.entity_id @@ -294,7 +265,7 @@ async def delete(self, records: Sequence[VectorDeletion]) -> None: ) await session.commit() - async def delete_entity(self, entity_id: int) -> None: + async def delete_entity(self, project_id: int, entity_id: int) -> None: await self.initialize() async with db.scoped_session(self._session_maker) as session: await session.execute( @@ -303,12 +274,12 @@ async def delete_entity(self, entity_id: int) -> None: "SELECT id FROM search_vector_chunks " "WHERE project_id = :project_id AND entity_id = :entity_id)" ), - {"project_id": self.scope.project_id, "entity_id": entity_id}, + {"project_id": project_id, "entity_id": entity_id}, ) await session.commit() - async def delete_orphans(self, _live_keys: Sequence[VectorKey]) -> None: - """Remove pgvector rows absent from the current ready manifest scope.""" + async def delete_orphans(self, project_id: int, _live_keys: Sequence[VectorKey]) -> None: + """Remove pgvector rows absent from one project's current ready manifest.""" await self.initialize() async with db.scoped_session(self._session_maker) as session: await session.execute( @@ -324,7 +295,7 @@ async def delete_orphans(self, _live_keys: Sequence[VectorKey]) -> None: "AND chunks.embedding_status = 'ready')" ), { - "project_id": self.scope.project_id, + "project_id": project_id, "embedding_identity": self.scope.embedding_identity, }, ) @@ -335,11 +306,21 @@ async def search( query: Sequence[float], *, limit: int, + projects: ProjectScope, ) -> list[VectorMatch]: - if not query or limit <= 0: + if not query or limit <= 0 or projects.is_empty: return [] validate_query_dimensions(self.scope, query) await self.initialize() + params: dict[str, object] = { + "query": self._format_vector(query), + "dimensions": self.scope.dimensions, + "embedding_identity": self.scope.embedding_identity, + "limit": limit, + } + # The scope binds its ids once; both predicates reference the same names. + embeddings_in_scope = projects.predicate("e.project_id", params) + chunks_in_scope = projects.predicate("c.project_id", params) async with db.scoped_session(self._session_maker) as session: result = await session.execute( text( @@ -347,9 +328,9 @@ async def search( "1 - (e.embedding <=> CAST(:query AS vector)) AS similarity " "FROM search_vector_embeddings e " "JOIN search_vector_chunks c ON c.id = e.chunk_id " - "WHERE e.project_id = :project_id " + f"WHERE {embeddings_in_scope} " "AND e.embedding_dims = :dimensions " - "AND c.project_id = :project_id " + f"AND {chunks_in_scope} " "AND c.vector_index = 'pgvector' " "AND c.embedding_status = 'ready' " "AND c.embedding_model = :embedding_identity " @@ -358,13 +339,7 @@ async def search( "c.entity_id ASC, c.chunk_key ASC " "LIMIT :limit" ), - { - "query": self._format_vector(query), - "project_id": self.scope.project_id, - "dimensions": self.scope.dimensions, - "embedding_identity": self.scope.embedding_identity, - "limit": limit, - }, + params, ) return [ VectorMatch( diff --git a/src/basic_memory/repository/postgres_search_repository.py b/src/basic_memory/repository/postgres_search_repository.py index 06f4e57d8..b57b4eb01 100644 --- a/src/basic_memory/repository/postgres_search_repository.py +++ b/src/basic_memory/repository/postgres_search_repository.py @@ -113,11 +113,7 @@ def __init__( ) vector_index = PgVectorIndex( session_maker, - build_vector_index_scope( - self._app_config, - self._embedding_provider, - project_id, - ), + build_vector_index_scope(self._app_config, self._embedding_provider), ) self._semantic_vector_index_name = effective_name self._semantic_vector_index = vector_index @@ -314,11 +310,7 @@ async def _ensure_vector_tables(self) -> None: self._semantic_vector_index_name = "pgvector" self._semantic_vector_index = PgVectorIndex( self.session_maker, - build_vector_index_scope( - self._app_config, - self._embedding_provider, - self.project_id, - ), + build_vector_index_scope(self._app_config, self._embedding_provider), ) if self._vector_tables_initialized: return diff --git a/src/basic_memory/repository/search_reader.py b/src/basic_memory/repository/search_reader.py index d5f925d54..0b5c81eb5 100644 --- a/src/basic_memory/repository/search_reader.py +++ b/src/basic_memory/repository/search_reader.py @@ -260,7 +260,9 @@ async def _run_vector_query( return [] if not self.vector.external: - matches = await self.vector.index.search(query_embedding, limit=candidate_limit) + matches = await self.vector.index.search( + query_embedding, limit=candidate_limit, projects=self.scope + ) if trace is not None: trace.readiness = await read_manifest_readiness( session, @@ -272,7 +274,9 @@ async def _run_vector_query( scan_limit = min(candidate_limit, VECTOR_FILTER_SCAN_LIMIT) while True: - matches = await self.vector.index.search(query_embedding, limit=scan_limit) + matches = await self.vector.index.search( + query_embedding, limit=scan_limit, projects=self.scope + ) if trace is not None and trace.readiness is None: trace.readiness = await read_manifest_readiness( session, diff --git a/src/basic_memory/repository/search_repository.py b/src/basic_memory/repository/search_repository.py index 2b9828202..b21aca5b1 100644 --- a/src/basic_memory/repository/search_repository.py +++ b/src/basic_memory/repository/search_repository.py @@ -235,7 +235,6 @@ def create_search_repository( embedding_provider = create_embedding_provider(config) vector_index_name, vector_index = create_semantic_vector_index( session_maker=session_maker, - project_id=project_id, app_config=config, database_backend=database_backend, embedding_provider=embedding_provider, diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 15e2dd506..8935b307d 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -621,7 +621,7 @@ async def _persist_embeddings( # the authoritative SQL database. Hold its manifest lock across # adapter I/O so a newer prepare cannot advance this generation # before the external write and ready transition complete. - await self._semantic_vector_index.upsert(records) + await self._semantic_vector_index.upsert(self.project_id, records) await self._mark_embedding_jobs_ready( session, params=params, @@ -632,7 +632,7 @@ async def _persist_embeddings( # Built-in adapters share the authoritative database. They verify and lock # each record's source_hash inside the same transaction as their vector write. - await self._semantic_vector_index.upsert(records) + await self._semantic_vector_index.upsert(self.project_id, records) async with db.scoped_session(self.session_maker) as session: await self._mark_embedding_jobs_ready( session, @@ -834,7 +834,7 @@ async def _finalize_prepared_vector_deletions( ] if not current_deletions: return - await self._semantic_vector_index.delete(current_deletions) + await self._semantic_vector_index.delete(self.project_id, current_deletions) await session.execute( text( "DELETE FROM search_vector_chunks " @@ -846,7 +846,7 @@ async def _finalize_prepared_vector_deletions( await session.commit() return - await self._semantic_vector_index.delete(deletions) + await self._semantic_vector_index.delete(self.project_id, deletions) if self._semantic_vector_index_name in BUILT_IN_VECTOR_INDEX_NAMES: return async with db.scoped_session(self.session_maker) as session: @@ -1209,7 +1209,7 @@ async def _delete_external_entity_vectors_locked( await self._semantic_vector_index.initialize() for entity_id in deleted_entity_ids: - await self._semantic_vector_index.delete_entity(entity_id) + await self._semantic_vector_index.delete_entity(self.project_id, entity_id) async def _delete_project_builtin_vector_rows(self, session: AsyncSession) -> None: """Delete backend-owned vector rows before their SQL manifest is removed.""" @@ -1479,11 +1479,11 @@ async def reconcile_vector_index(self) -> None: ] if external_vector_index: - await self._semantic_vector_index.delete_orphans(live_keys) + await self._semantic_vector_index.delete_orphans(self.project_id, live_keys) await session.commit() return - await self._semantic_vector_index.delete_orphans(live_keys) + await self._semantic_vector_index.delete_orphans(self.project_id, live_keys) # ------------------------------------------------------------------ # Shared semantic search: guard, text processing, chunking diff --git a/src/basic_memory/repository/semantic_vector_index.py b/src/basic_memory/repository/semantic_vector_index.py index e869fb0a8..bb6b5a2b7 100644 --- a/src/basic_memory/repository/semantic_vector_index.py +++ b/src/basic_memory/repository/semantic_vector_index.py @@ -6,25 +6,28 @@ from dataclasses import dataclass from typing import Protocol, runtime_checkable +from basic_memory.repository.search_scope import ProjectScope + @dataclass(frozen=True, slots=True) class VectorIndexScope: - """Stable project storage identity plus the current embedding schema.""" + """The database namespace and embedding schema an adapter stores vectors under. + + Projects are partitions inside it: every write names the project it touches and a + search names the projects it reads, so one adapter serves a whole database. + """ namespace: str - project_id: int embedding_identity: str dimensions: int - @property - def storage_key(self) -> tuple[str, int]: - """Return the stable isolation key external adapters must use for storage.""" - return (self.namespace, self.project_id) - @dataclass(frozen=True, slots=True) class VectorKey: - """Backend-independent identity for one semantic chunk vector.""" + """Backend-independent identity for one semantic chunk vector. + + Entity ids are database-wide primary keys, so the pair is unique across projects. + """ entity_id: int chunk_key: str @@ -68,19 +71,19 @@ class SemanticVectorIndex(Protocol): def scope(self) -> VectorIndexScope: ... async def initialize(self) -> None: - """Create or validate backend storage for the configured scope.""" + """Create or validate backend storage shared by every project in the scope.""" ... - async def upsert(self, records: Sequence[VectorRecord]) -> None: - """Insert or replace vectors only for each record's source generation.""" + async def upsert(self, project_id: int, records: Sequence[VectorRecord]) -> None: + """Insert or replace one project's vectors, only for each record's source generation.""" ... - async def delete(self, records: Sequence[VectorDeletion]) -> None: - """Delete vectors only for each record's source generation.""" + async def delete(self, project_id: int, records: Sequence[VectorDeletion]) -> None: + """Delete one project's vectors, only for each record's source generation.""" ... - async def delete_entity(self, entity_id: int) -> None: - """Delete every vector owned by an entity in this scope.""" + async def delete_entity(self, project_id: int, entity_id: int) -> None: + """Delete every vector owned by an entity in one project.""" ... async def search( @@ -88,8 +91,9 @@ async def search( query: Sequence[float], *, limit: int, + projects: ProjectScope, ) -> list[VectorMatch]: - """Return nearest matches ordered by normalized cosine similarity.""" + """Return nearest matches within ``projects``, ordered by normalized similarity.""" ... @@ -100,8 +104,8 @@ class SemanticVectorIndexReconciler(Protocol): @property def scope(self) -> VectorIndexScope: ... - async def delete_orphans(self, live_keys: Sequence[VectorKey]) -> None: - """Delete scoped vectors whose stable keys are not in ``live_keys``.""" + async def delete_orphans(self, project_id: int, live_keys: Sequence[VectorKey]) -> None: + """Delete one project's vectors whose stable keys are not in ``live_keys``.""" ... diff --git a/src/basic_memory/repository/semantic_vector_index_factory.py b/src/basic_memory/repository/semantic_vector_index_factory.py index d9bd91869..f29d00517 100644 --- a/src/basic_memory/repository/semantic_vector_index_factory.py +++ b/src/basic_memory/repository/semantic_vector_index_factory.py @@ -76,12 +76,13 @@ def _database_namespace(app_config: BasicMemoryConfig) -> str: def build_vector_index_scope( app_config: BasicMemoryConfig, provider: EmbeddingProvider, - project_id: int, ) -> VectorIndexScope: - """Build the explicit isolation contract handed to every vector adapter.""" + """Build the storage identity handed to every vector adapter. + + The database namespace and embedding schema; projects are named per operation. + """ return VectorIndexScope( namespace=_database_namespace(app_config), - project_id=project_id, embedding_identity=semantic_embedding_identity(provider), dimensions=provider.dimensions, ) @@ -109,14 +110,13 @@ def _create_milvus_index( def create_semantic_vector_index( *, session_maker: async_sessionmaker[AsyncSession], - project_id: int, app_config: BasicMemoryConfig, database_backend: DatabaseBackend, embedding_provider: EmbeddingProvider, ) -> tuple[str, SemanticVectorIndex]: """Create the vector adapter selected by the validated application config.""" name = resolve_semantic_vector_index_name(app_config, database_backend) - scope = build_vector_index_scope(app_config, embedding_provider, project_id) + scope = build_vector_index_scope(app_config, embedding_provider) if name == "sqlite-vec": from basic_memory.repository.sqlite_vec_index import SQLiteVecIndex diff --git a/src/basic_memory/repository/sqlite_search_repository.py b/src/basic_memory/repository/sqlite_search_repository.py index 0861023b7..7f50b8df5 100644 --- a/src/basic_memory/repository/sqlite_search_repository.py +++ b/src/basic_memory/repository/sqlite_search_repository.py @@ -85,11 +85,7 @@ def __init__( self._vector_dimensions = self._embedding_provider.dimensions self._semantic_vector_index = vector_index or SQLiteVecIndex( session_maker, - build_vector_index_scope( - self._app_config, - self._embedding_provider, - project_id, - ), + build_vector_index_scope(self._app_config, self._embedding_provider), ) @override @@ -269,11 +265,7 @@ async def _ensure_vector_tables(self) -> None: assert self._embedding_provider is not None self._semantic_vector_index = SQLiteVecIndex( self.session_maker, - build_vector_index_scope( - self._app_config, - self._embedding_provider, - self.project_id, - ), + build_vector_index_scope(self._app_config, self._embedding_provider), ) if self._vector_tables_initialized: return diff --git a/src/basic_memory/repository/sqlite_vec_index.py b/src/basic_memory/repository/sqlite_vec_index.py index c3344c31c..f85f40248 100644 --- a/src/basic_memory/repository/sqlite_vec_index.py +++ b/src/basic_memory/repository/sqlite_vec_index.py @@ -13,6 +13,7 @@ from basic_memory import db from basic_memory.models.search import create_sqlite_search_vector_embeddings +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError from basic_memory.repository.semantic_vector_index import ( VectorDeletion, @@ -126,43 +127,14 @@ async def initialize(self) -> None: await session.commit() self._initialized = True - async def _rowids_by_key( - self, - session: AsyncSession, - keys: Sequence[VectorKey], - ) -> dict[VectorKey, int]: - if not keys: - return {} - params: dict[str, object] = {"project_id": self.scope.project_id} - predicates: list[str] = [] - for index, key in enumerate(keys): - params[f"entity_id_{index}"] = key.entity_id - params[f"chunk_key_{index}"] = key.chunk_key - predicates.append( - f"(entity_id = :entity_id_{index} AND chunk_key = :chunk_key_{index})" - ) - result = await session.execute( - text( - "SELECT id, entity_id, chunk_key FROM search_vector_chunks " - "WHERE project_id = :project_id AND (" + " OR ".join(predicates) + ")" - ), - params, - ) - return { - VectorKey(entity_id=int(row["entity_id"]), chunk_key=str(row["chunk_key"])): int( - row["id"] - ) - for row in result.mappings().all() - } - - async def upsert(self, records: Sequence[VectorRecord]) -> None: + async def upsert(self, project_id: int, records: Sequence[VectorRecord]) -> None: if not records: return validate_vector_dimensions(self.scope, records) await self.initialize() async with db.scoped_session(self._session_maker) as session: await self._ensure_loaded(session) - params: dict[str, object] = {"project_id": self.scope.project_id} + params: dict[str, object] = {"project_id": project_id} predicates: list[str] = [] records_by_key = {record.key: record for record in records} for index, record in enumerate(records): @@ -223,13 +195,13 @@ async def upsert(self, records: Sequence[VectorRecord]) -> None: ) await session.commit() - async def delete(self, records: Sequence[VectorDeletion]) -> None: + async def delete(self, project_id: int, records: Sequence[VectorDeletion]) -> None: if not records: return await self.initialize() async with db.scoped_session(self._session_maker) as session: await self._ensure_loaded(session) - params: dict[str, object] = {"project_id": self.scope.project_id} + params: dict[str, object] = {"project_id": project_id} predicates: list[str] = [] for index, record in enumerate(records): params[f"entity_id_{index}"] = record.key.entity_id @@ -264,7 +236,7 @@ async def delete(self, records: Sequence[VectorDeletion]) -> None: ) await session.commit() - async def delete_entity(self, entity_id: int) -> None: + async def delete_entity(self, project_id: int, entity_id: int) -> None: await self.initialize() async with db.scoped_session(self._session_maker) as session: await self._ensure_loaded(session) @@ -274,12 +246,12 @@ async def delete_entity(self, entity_id: int) -> None: "SELECT id FROM search_vector_chunks " "WHERE project_id = :project_id AND entity_id = :entity_id)" ), - {"project_id": self.scope.project_id, "entity_id": entity_id}, + {"project_id": project_id, "entity_id": entity_id}, ) await session.commit() - async def delete_orphans(self, _live_keys: Sequence[VectorKey]) -> None: - """Remove sqlite-vec rows absent from the current ready manifest scope.""" + async def delete_orphans(self, project_id: int, _live_keys: Sequence[VectorKey]) -> None: + """Remove sqlite-vec rows absent from one project's current ready manifest.""" await self.initialize() async with db.scoped_session(self._session_maker) as session: await self._ensure_loaded(session) @@ -316,7 +288,7 @@ async def delete_orphans(self, _live_keys: Sequence[VectorKey]) -> None: "AND embedding_status = 'ready'))" ), { - "project_id": self.scope.project_id, + "project_id": project_id, "embedding_identity": self.scope.embedding_identity, }, ) @@ -327,12 +299,20 @@ async def search( query: Sequence[float], *, limit: int, + projects: ProjectScope, ) -> list[VectorMatch]: - if not query or limit <= 0: + if not query or limit <= 0 or projects.is_empty: return [] validate_query_dimensions(self.scope, query) await self.initialize() vector_k = min(limit, SQLITE_VEC_MAX_K) + params: dict[str, object] = { + "query": json.dumps(list(query)), + "vector_k": vector_k, + "embedding_identity": self.scope.embedding_identity, + "limit": limit, + } + chunks_in_scope = projects.predicate("c.project_id", params) async with db.scoped_session(self._session_maker) as session: await self._ensure_loaded(session) result = await session.execute( @@ -345,20 +325,14 @@ async def search( "FROM vector_matches " "JOIN search_vector_chunks c ON c.id = vector_matches.rowid " "AND c.source_hash = vector_matches.source_hash " - "WHERE c.project_id = :project_id " + f"WHERE {chunks_in_scope} " "AND c.vector_index = 'sqlite-vec' " "AND c.embedding_status = 'ready' " "AND c.embedding_model = :embedding_identity " "ORDER BY vector_matches.distance ASC, " "c.entity_id ASC, c.chunk_key ASC LIMIT :limit" ), - { - "query": json.dumps(list(query)), - "vector_k": vector_k, - "project_id": self.scope.project_id, - "embedding_identity": self.scope.embedding_identity, - "limit": limit, - }, + params, ) return [ VectorMatch( diff --git a/test-int/semantic/conftest.py b/test-int/semantic/conftest.py index c5c87559a..53ea89946 100644 --- a/test-int/semantic/conftest.py +++ b/test-int/semantic/conftest.py @@ -293,7 +293,6 @@ async def create_search_service( if embedding_provider is not None: vector_index_name, vector_index = create_semantic_vector_index( session_maker=session_maker, - project_id=project.id, app_config=app_config, database_backend=combo.backend, embedding_provider=embedding_provider, diff --git a/test-int/semantic/test_milvus_lite.py b/test-int/semantic/test_milvus_lite.py index 7f4db5e08..a637204c6 100644 --- a/test-int/semantic/test_milvus_lite.py +++ b/test-int/semantic/test_milvus_lite.py @@ -11,6 +11,7 @@ from basic_memory.repository.milvus_config import MilvusSettings from basic_memory.repository.milvus_index import MilvusVectorIndex +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_vector_index import ( VectorDeletion, VectorIndexScope, @@ -23,6 +24,9 @@ pytest.mark.skipif(sys.platform == "win32", reason="Milvus Lite does not support Windows"), ] +PROJECT = 7 +PROJECTS = ProjectScope.single(PROJECT) + _RESTART_SCRIPT = """ import asyncio @@ -30,6 +34,7 @@ from basic_memory.repository.milvus_config import MilvusSettings from basic_memory.repository.milvus_index import MilvusVectorIndex +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_vector_index import ( VectorIndexScope, VectorKey, @@ -41,7 +46,6 @@ async def main() -> None: phase, database_path = sys.argv[1:] scope = VectorIndexScope( namespace="restart-database", - project_id=7, embedding_identity="Stub:model", dimensions=3, ) @@ -50,18 +54,19 @@ async def main() -> None: if phase == "write": await index.upsert( + 7, [ VectorRecord( key=key, source_hash="auth-v1", values=(1.0, 0.0, 0.0), ) - ] + ], ) return - await index.delete_orphans([key]) - matches = await index.search((1.0, 0.0, 0.0), limit=1) + await index.delete_orphans(7, [key]) + matches = await index.search((1.0, 0.0, 0.0), limit=1, projects=ProjectScope.single(7)) assert [match.key for match in matches] == [key] @@ -83,7 +88,6 @@ def _run_restart_phase(phase: str, database_path: str) -> None: async def test_milvus_lite_vector_lifecycle(tmp_path) -> None: scope = VectorIndexScope( namespace="integration-database", - project_id=7, embedding_identity="Stub:model", dimensions=3, ) @@ -94,6 +98,7 @@ async def test_milvus_lite_vector_lifecycle(tmp_path) -> None: auth_key = VectorKey(entity_id=1, chunk_key="summary:0") database_key = VectorKey(entity_id=2, chunk_key="summary:0") await index.upsert( + PROJECT, [ VectorRecord( key=auth_key, @@ -105,21 +110,23 @@ async def test_milvus_lite_vector_lifecycle(tmp_path) -> None: source_hash="database-v1", values=(0.0, 1.0, 0.0), ), - ] + ], ) - matches = await index.search((1.0, 0.0, 0.0), limit=2) + matches = await index.search((1.0, 0.0, 0.0), limit=2, projects=PROJECTS) assert matches[0].key == auth_key assert matches[0].similarity == pytest.approx(1.0) - await index.delete([VectorDeletion(key=auth_key, source_hash="stale-generation")]) - assert (await index.search((1.0, 0.0, 0.0), limit=2))[0].key == auth_key + await index.delete(PROJECT, [VectorDeletion(key=auth_key, source_hash="stale-generation")]) + assert (await index.search((1.0, 0.0, 0.0), limit=2, projects=PROJECTS))[0].key == auth_key - await index.delete([VectorDeletion(key=auth_key, source_hash="auth-v1")]) - assert [match.key for match in await index.search((1.0, 0.0, 0.0), limit=2)] == [database_key] + await index.delete(PROJECT, [VectorDeletion(key=auth_key, source_hash="auth-v1")]) + assert [ + match.key for match in await index.search((1.0, 0.0, 0.0), limit=2, projects=PROJECTS) + ] == [database_key] - await index.delete_orphans([]) - assert await index.search((1.0, 0.0, 0.0), limit=2) == [] + await index.delete_orphans(PROJECT, []) + assert await index.search((1.0, 0.0, 0.0), limit=2, projects=PROJECTS) == [] def test_milvus_lite_reloads_collection_after_process_restart(tmp_path) -> None: diff --git a/test-int/test_embedding_status_vec0.py b/test-int/test_embedding_status_vec0.py index 1ccbee5cc..d4c3ee949 100644 --- a/test-int/test_embedding_status_vec0.py +++ b/test-int/test_embedding_status_vec0.py @@ -137,13 +137,14 @@ async def test_embedding_status_reads_real_vec0_table(engine_factory, test_proje # An obsolete embedding result must not claim the stable vec0 row after the # manifest has advanced to a newer source generation. await search_repo._semantic_vector_index.upsert( + project_id, [ VectorRecord( key=VectorKey(entity_id=entity_id, chunk_key="chunk-1"), source_hash="stale-hash", values=tuple(_unit_vector(dimensions)), ) - ] + ], ) async with db.scoped_session(session_maker) as session: # sqlite-vec is loaded per connection. Windows may hand this assertion a @@ -153,13 +154,14 @@ async def test_embedding_status_reads_real_vec0_table(engine_factory, test_proje assert stale_count.scalar_one() == 0 await search_repo._semantic_vector_index.upsert( + project_id, [ VectorRecord( key=VectorKey(entity_id=entity_id, chunk_key="chunk-1"), source_hash="hash", values=tuple(_unit_vector(dimensions)), ) - ] + ], ) async with db.scoped_session(session_maker) as session: diff --git a/tests/repository/test_milvus_index.py b/tests/repository/test_milvus_index.py index 47790c72d..0dfe584c5 100644 --- a/tests/repository/test_milvus_index.py +++ b/tests/repository/test_milvus_index.py @@ -23,6 +23,7 @@ MilvusVectorIndex, collection_name, ) +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_vector_index import ( SemanticVectorIndex, SemanticVectorIndexReconciler, @@ -35,6 +36,11 @@ create_semantic_vector_index, ) +PROJECT = 42 +OTHER_PROJECT = 43 +PROJECTS = ProjectScope.single(PROJECT) +QUERY = [1.0, 0.0, 0.0] + class FakeRepository: """In-memory recorder for the blocking Milvus repository.""" @@ -57,6 +63,8 @@ def __init__( self.ids: list[str] = [] self.id_deletes: list[tuple[str, list[str]]] = [] self.matches: list[MilvusStoredMatch] = [] + # Per-collection answers; ``matches`` is the answer for any collection not listed. + self.matches_by_collection: dict[str, list[MilvusStoredMatch]] = {} self.searches: list[tuple[str, list[float], int]] = [] self.closed = 0 @@ -103,7 +111,7 @@ def search( limit: int, ) -> list[MilvusStoredMatch]: self.searches.append((collection_name, list(query), limit)) - return self.matches + return self.matches_by_collection.get(collection_name, self.matches) def close(self) -> None: self.closed += 1 @@ -172,7 +180,6 @@ def runtime_log_attrs(self) -> dict[str, Any]: def scope() -> VectorIndexScope: return VectorIndexScope( namespace="basic-memory-database", - project_id=42, embedding_identity="Provider:model-a", dimensions=3, ) @@ -195,84 +202,99 @@ def _index( ) -def test_collection_name_uses_only_stable_scope_identity( +def test_collection_name_uses_only_stable_project_identity( scope: VectorIndexScope, settings: MilvusSettings, ) -> None: changed_schema = VectorIndexScope( namespace=scope.namespace, - project_id=scope.project_id, embedding_identity="Provider:model-b", dimensions=9, ) - other_project = VectorIndexScope( - namespace=scope.namespace, - project_id=43, - embedding_identity=scope.embedding_identity, - dimensions=scope.dimensions, + + assert collection_name(settings, scope, PROJECT) == collection_name( + settings, changed_schema, PROJECT ) + assert collection_name(settings, scope, PROJECT) != collection_name( + settings, scope, OTHER_PROJECT + ) + assert collection_name(settings, scope, PROJECT).startswith("basic_memory_") + + +# --- Collection validation happens on a project's first use --- + + +@pytest.mark.asyncio +async def test_initialize_prepares_nothing_shared( + scope: VectorIndexScope, + settings: MilvusSettings, +) -> None: + """Collections are per project, so the database-wide hook has nothing to create.""" + repository = FakeRepository() - assert collection_name(settings, scope) == collection_name(settings, changed_schema) - assert collection_name(settings, scope) != collection_name(settings, other_project) - assert collection_name(settings, scope).startswith("basic_memory_") + await _index(scope, settings, repository).initialize() + + assert repository.created == [] + assert repository.closed == 0 @pytest.mark.asyncio -async def test_initialize_creates_missing_collection_once( +async def test_first_use_creates_missing_collection_once( scope: VectorIndexScope, settings: MilvusSettings, ) -> None: repository = FakeRepository() index = _index(scope, settings, repository) - await index.initialize() - await index.initialize() + await index.search(QUERY, limit=1, projects=PROJECTS) + await index.search(QUERY, limit=1, projects=PROJECTS) - assert repository.created == [(collection_name(settings, scope), scope.dimensions)] - assert repository.closed == 1 + assert repository.created == [(collection_name(settings, scope, PROJECT), scope.dimensions)] + # One validation plus one search per call. + assert repository.closed == 3 @pytest.mark.asyncio -async def test_initialize_accepts_compatible_collection_create_race( +async def test_first_use_accepts_compatible_collection_create_race( scope: VectorIndexScope, settings: MilvusSettings, ) -> None: repository = FakeRepository(create_result=False, race_dimensions=scope.dimensions) - await _index(scope, settings, repository).initialize() + await _index(scope, settings, repository).search(QUERY, limit=1, projects=PROJECTS) - assert repository.created == [(collection_name(settings, scope), scope.dimensions)] + assert repository.created == [(collection_name(settings, scope, PROJECT), scope.dimensions)] assert repository.dimensions == scope.dimensions - assert repository.loaded == [collection_name(settings, scope)] + assert repository.loaded == [collection_name(settings, scope, PROJECT)] @pytest.mark.asyncio -async def test_initialize_rejects_incompatible_collection_create_race( +async def test_first_use_rejects_incompatible_collection_create_race( scope: VectorIndexScope, settings: MilvusSettings, ) -> None: repository = FakeRepository(create_result=False, race_dimensions=99) with pytest.raises(RuntimeError, match="Refusing to replace shared vector storage"): - await _index(scope, settings, repository).initialize() + await _index(scope, settings, repository).search(QUERY, limit=1, projects=PROJECTS) assert repository.dimensions == 99 assert repository.loaded == [] @pytest.mark.asyncio -async def test_initialize_rejects_collection_disappearing_after_create_race( +async def test_first_use_rejects_collection_disappearing_after_create_race( scope: VectorIndexScope, settings: MilvusSettings, ) -> None: repository = FakeRepository(create_result=False) with pytest.raises(RuntimeError, match="disappeared after a concurrent create"): - await _index(scope, settings, repository).initialize() + await _index(scope, settings, repository).search(QUERY, limit=1, projects=PROJECTS) @pytest.mark.asyncio -async def test_initialize_preserves_collection_on_dimension_mismatch( +async def test_first_use_preserves_collection_on_dimension_mismatch( scope: VectorIndexScope, settings: MilvusSettings, ) -> None: @@ -280,7 +302,7 @@ async def test_initialize_preserves_collection_on_dimension_mismatch( index = _index(scope, settings, repository) with pytest.raises(RuntimeError, match="Refusing to replace shared vector storage"): - await index.initialize() + await index.search(QUERY, limit=1, projects=PROJECTS) assert repository.created == [] assert repository.dimensions == 99 @@ -288,36 +310,39 @@ async def test_initialize_preserves_collection_on_dimension_mismatch( @pytest.mark.asyncio -async def test_initialize_accepts_matching_collection( +async def test_first_use_accepts_matching_collection( scope: VectorIndexScope, settings: MilvusSettings, ) -> None: repository = FakeRepository(dimensions=scope.dimensions) - await _index(scope, settings, repository).initialize() + await _index(scope, settings, repository).search(QUERY, limit=1, projects=PROJECTS) assert repository.created == [] - assert repository.loaded == [collection_name(settings, scope)] + assert repository.loaded == [collection_name(settings, scope, PROJECT)] @pytest.mark.asyncio -async def test_concurrent_initialize_rechecks_state_inside_lock( +async def test_concurrent_first_use_rechecks_state_inside_lock( scope: VectorIndexScope, settings: MilvusSettings, ) -> None: repository = FakeRepository() index = _index(scope, settings, repository) - await index._initialize_lock.acquire() - waiting_initialize = asyncio.create_task(index.initialize()) + await index._collection_lock.acquire() + waiting = asyncio.create_task(index._ensure_collection(PROJECT)) await asyncio.sleep(0) - index._initialized = True - index._initialize_lock.release() - await waiting_initialize + index._ready_projects.add(PROJECT) + index._collection_lock.release() + assert await waiting == collection_name(settings, scope, PROJECT) assert repository.closed == 0 +# --- Writes --- + + @pytest.mark.asyncio async def test_upsert_preserves_stable_key_generation_and_values( scope: VectorIndexScope, @@ -331,9 +356,10 @@ async def test_upsert_preserves_stable_key_generation_and_values( values=(1.0, 0.0, -1.0), ) - await index.upsert([record]) + await index.upsert(PROJECT, [record]) - _, stored_records = repository.upserts[0] + collection, stored_records = repository.upserts[0] + assert collection == collection_name(settings, scope, PROJECT) assert len(stored_records) == 1 assert stored_records[0].entity_id == 7 assert stored_records[0].chunk_key == "summary:0" @@ -352,13 +378,14 @@ async def test_upsert_rejects_wrong_dimensions_before_milvus_call( with pytest.raises(ValueError, match="expected 3, got 2"): await index.upsert( + PROJECT, [ VectorRecord( key=VectorKey(entity_id=7, chunk_key="summary:0"), source_hash="source-a", values=(1.0, 0.0), ) - ] + ], ) assert repository.upserts == [] @@ -378,14 +405,14 @@ async def test_mutations_finish_before_propagating_cancellation( if operation == "upsert": mutation = index.upsert( - [VectorRecord(key=key, source_hash="source-a", values=(1.0, 0.0, 0.0))] + PROJECT, [VectorRecord(key=key, source_hash="source-a", values=(1.0, 0.0, 0.0))] ) elif operation == "delete": - mutation = index.delete([VectorDeletion(key=key, source_hash="source-a")]) + mutation = index.delete(PROJECT, [VectorDeletion(key=key, source_hash="source-a")]) elif operation == "delete_entity": - mutation = index.delete_entity(key.entity_id) + mutation = index.delete_entity(PROJECT, key.entity_id) else: - mutation = index.delete_orphans([]) + mutation = index.delete_orphans(PROJECT, []) mutation_task = asyncio.create_task(mutation) async with asyncio.timeout(2): @@ -415,24 +442,29 @@ async def test_delete_forwards_source_generation( index = _index(scope, settings, repository) key = VectorKey(entity_id=7, chunk_key="summary:0") - await index.delete([VectorDeletion(key=key, source_hash="source-a")]) + await index.delete(PROJECT, [VectorDeletion(key=key, source_hash="source-a")]) - _, deletions = repository.record_deletes[0] + collection, deletions = repository.record_deletes[0] + assert collection == collection_name(settings, scope, PROJECT) assert deletions[0][1] == "source-a" assert len(deletions[0][0]) == 64 @pytest.mark.asyncio -async def test_delete_entity_uses_project_collection( +async def test_delete_entity_uses_the_named_project_collection( scope: VectorIndexScope, settings: MilvusSettings, ) -> None: repository = FakeRepository(dimensions=scope.dimensions) index = _index(scope, settings, repository) - await index.delete_entity(77) + await index.delete_entity(PROJECT, 77) + await index.delete_entity(OTHER_PROJECT, 78) - assert repository.entity_deletes == [(collection_name(settings, scope), 77)] + assert repository.entity_deletes == [ + (collection_name(settings, scope, PROJECT), 77), + (collection_name(settings, scope, OTHER_PROJECT), 78), + ] @pytest.mark.asyncio @@ -446,18 +478,19 @@ async def test_reconciliation_deletes_only_absent_stable_keys( stale_key = VectorKey(entity_id=2, chunk_key="stale") await index.upsert( + PROJECT, [ VectorRecord(key=live_key, source_hash="a", values=(1.0, 0.0, 0.0)), VectorRecord(key=stale_key, source_hash="b", values=(0.0, 1.0, 0.0)), - ] + ], ) _, stored_records = repository.upserts[0] repository.ids = [record.record_id for record in stored_records] - await index.delete_orphans([live_key]) + await index.delete_orphans(PROJECT, [live_key]) assert repository.id_deletes == [ - (collection_name(settings, scope), [stored_records[1].record_id]) + (collection_name(settings, scope, PROJECT), [stored_records[1].record_id]) ] @@ -469,7 +502,7 @@ async def test_reconciliation_is_noop_without_orphans( repository = FakeRepository(dimensions=scope.dimensions) index = _index(scope, settings, repository) - await index.delete_orphans([]) + await index.delete_orphans(PROJECT, []) assert repository.id_deletes == [] @@ -483,11 +516,14 @@ async def test_reconciliation_deletes_orphans_incrementally( repository.ids = [f"orphan-{index}" for index in range(600)] index = _index(scope, settings, repository) - await index.delete_orphans([]) + await index.delete_orphans(PROJECT, []) assert [len(record_ids) for _, record_ids in repository.id_deletes] == [256, 256, 88] +# --- Search --- + + @pytest.mark.asyncio async def test_search_clamps_milvus_cosine_scores_and_orders_ties( scope: VectorIndexScope, @@ -502,25 +538,60 @@ async def test_search_clamps_milvus_cosine_scores_and_orders_ties( ] index = _index(scope, settings, repository) - matches = await index.search([1.0, 0.0, 0.0], limit=4) + matches = await index.search(QUERY, limit=4, projects=PROJECTS) assert [match.similarity for match in matches] == [1.0, 1.0, 0.1, 0.0] assert [match.key.entity_id for match in matches] == [0, 1, 2, 3] - assert repository.searches == [(collection_name(settings, scope), [1.0, 0.0, 0.0], 4)] + assert repository.searches == [(collection_name(settings, scope, PROJECT), QUERY, 4)] + + +@pytest.mark.asyncio +async def test_search_merges_the_collections_in_scope( + scope: VectorIndexScope, + settings: MilvusSettings, +) -> None: + """Milvus has no cross-collection search, so a wider scope merges per-project answers.""" + repository = FakeRepository(dimensions=scope.dimensions) + first = collection_name(settings, scope, PROJECT) + second = collection_name(settings, scope, OTHER_PROJECT) + repository.matches_by_collection = { + first: [ + MilvusStoredMatch(entity_id=1, chunk_key="a", score=0.9), + MilvusStoredMatch(entity_id=2, chunk_key="a", score=0.3), + ], + second: [ + MilvusStoredMatch(entity_id=3, chunk_key="a", score=0.8), + MilvusStoredMatch(entity_id=4, chunk_key="a", score=0.7), + ], + } + index = _index(scope, settings, repository) + + matches = await index.search(QUERY, limit=3, projects=ProjectScope.of([OTHER_PROJECT, PROJECT])) + + assert [(match.key.entity_id, match.similarity) for match in matches] == [ + (1, 0.9), + (3, 0.8), + (4, 0.7), + ] + # Each project's collection is asked for its own top ``limit`` before the merge. + assert repository.searches == [(first, QUERY, 3), (second, QUERY, 3)] + # Both collections were validated before being searched. + assert repository.loaded == [first, second] @pytest.mark.asyncio -async def test_empty_operations_do_not_initialize( +async def test_empty_operations_do_not_touch_milvus( scope: VectorIndexScope, settings: MilvusSettings, ) -> None: repository = FakeRepository() index = _index(scope, settings, repository) - await index.upsert([]) - await index.delete([]) - assert await index.search([], limit=10) == [] - assert await index.search([1.0, 0.0, 0.0], limit=0) == [] + await index.upsert(PROJECT, []) + await index.delete(PROJECT, []) + assert await index.search([], limit=10, projects=PROJECTS) == [] + assert await index.search(QUERY, limit=0, projects=PROJECTS) == [] + assert await index.search(QUERY, limit=10, projects=ProjectScope.of([])) == [] assert repository.closed == 0 @@ -535,7 +606,6 @@ def test_first_party_factory_loads_milvus() -> None: name, index = create_semantic_vector_index( session_maker=session_maker, - project_id=42, app_config=app_config, database_backend=DatabaseBackend.POSTGRES, embedding_provider=StubEmbeddingProvider(), @@ -545,4 +615,4 @@ def test_first_party_factory_loads_milvus() -> None: assert isinstance(index, MilvusVectorIndex) assert isinstance(index, SemanticVectorIndex) assert isinstance(index, SemanticVectorIndexReconciler) - assert index.scope.project_id == 42 + assert index.scope.dimensions == 3 diff --git a/tests/repository/test_pgvector_index.py b/tests/repository/test_pgvector_index.py index e1fa8bfe8..c46abca42 100644 --- a/tests/repository/test_pgvector_index.py +++ b/tests/repository/test_pgvector_index.py @@ -10,6 +10,7 @@ from basic_memory.repository import pgvector_index as pgvector_index_module from basic_memory.repository.pgvector_index import PgVectorIndex +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError from basic_memory.repository.semantic_vector_index import ( VectorDeletion, @@ -93,10 +94,13 @@ async def commit(self) -> None: self.commit_count += 1 +PROJECT = 7 +PROJECTS = ProjectScope.single(PROJECT) + + def _scope(dimensions: int = 4) -> VectorIndexScope: return VectorIndexScope( namespace="basic-memory-test", - project_id=7, embedding_identity="stub:4", dimensions=dimensions, ) @@ -208,10 +212,11 @@ async def test_upsert_resolves_stable_keys_and_writes_one_batch(monkeypatch) -> index._initialized = True await index.upsert( + PROJECT, [ VectorRecord(key=key_a, source_hash="hash-a", values=(1.0, 0.0, 0.0, 0.0)), VectorRecord(key=key_b, source_hash="hash-b", values=(0.0, 1.0, 0.0, 0.0)), - ] + ], ) insert_call = next(call for call in session.calls if "INSERT INTO" in call[0]) @@ -246,7 +251,9 @@ async def test_upsert_skips_stale_source_generation(monkeypatch) -> None: index = PgVectorIndex(MagicMock(), _scope()) index._initialized = True - await index.upsert([VectorRecord(key=key, source_hash="old-hash", values=(1.0, 0.0, 0.0, 0.0))]) + await index.upsert( + PROJECT, [VectorRecord(key=key, source_hash="old-hash", values=(1.0, 0.0, 0.0, 0.0))] + ) lock_call = next(call for call in session.calls if "SELECT id, entity_id" in call[0]) assert "FOR UPDATE" in lock_call[0] @@ -263,7 +270,9 @@ async def test_upsert_rejects_missing_manifest_key(monkeypatch) -> None: index._initialized = True with pytest.raises(RuntimeError, match="manifest rows are missing"): - await index.upsert([VectorRecord(key=key, source_hash="hash", values=(1.0, 0.0, 0.0, 0.0))]) + await index.upsert( + PROJECT, [VectorRecord(key=key, source_hash="hash", values=(1.0, 0.0, 0.0, 0.0))] + ) @pytest.mark.asyncio @@ -274,10 +283,10 @@ async def test_delete_stable_keys_and_entity(monkeypatch) -> None: index = PgVectorIndex(MagicMock(), _scope()) index._initialized = True - await index.delete([]) - await index.delete([VectorDeletion(key=key, source_hash="hash")]) - await index.delete_entity(13) - await index.delete_orphans([key]) + await index.delete(PROJECT, []) + await index.delete(PROJECT, [VectorDeletion(key=key, source_hash="hash")]) + await index.delete_entity(PROJECT, 13) + await index.delete_orphans(PROJECT, [key]) delete_lock = next(call for call in session.calls if "SELECT id, entity_id" in call[0]) assert "source_hash = :source_hash_0" in delete_lock[0] @@ -307,12 +316,14 @@ async def test_search_returns_normalized_stable_matches(monkeypatch) -> None: index = PgVectorIndex(MagicMock(), _scope()) index._initialized = True - assert await index.search([], limit=5) == [] - assert await index.search([1.0, 0.0, 0.0, 0.0], limit=0) == [] + assert await index.search([], limit=5, projects=PROJECTS) == [] + assert await index.search([1.0, 0.0, 0.0, 0.0], limit=0, projects=PROJECTS) == [] + assert await index.search([1.0, 0.0, 0.0, 0.0], limit=5, projects=ProjectScope.of([])) == [] with pytest.raises(ValueError, match="expected 4, got 2"): - await index.search([1.0, 0.0], limit=5) + await index.search([1.0, 0.0], limit=5, projects=PROJECTS) + assert session.calls == [] - matches = await index.search([1.0, 0.0, 0.0, 0.0], limit=5) + matches = await index.search([1.0, 0.0, 0.0, 0.0], limit=5, projects=PROJECTS) assert [(match.key.entity_id, match.similarity) for match in matches] == [ (14, 1.0), @@ -322,8 +333,26 @@ async def test_search_returns_normalized_stable_matches(monkeypatch) -> None: assert "c.entity_id ASC, c.chunk_key ASC" in search_call[0] assert search_call[1] == { "query": "[1,0,0,0]", - "project_id": 7, + "scope_0": 7, "dimensions": 4, "embedding_identity": "stub:4", "limit": 5, } + + +@pytest.mark.asyncio +async def test_search_binds_every_project_in_scope(monkeypatch) -> None: + """A multi-project scope filters both the embedding and manifest rows by the same ids.""" + session = FakeSession(search_rows=[]) + _install_session(monkeypatch, session) + index = PgVectorIndex(MagicMock(), _scope()) + index._initialized = True + + await index.search([1.0, 0.0, 0.0, 0.0], limit=5, projects=ProjectScope.of([9, 7])) + + sql, params = next(call for call in session.calls if "AS similarity" in call[0]) + assert params is not None + assert "WHERE e.project_id IN (:scope_0, :scope_1)" in sql + assert "AND c.project_id IN (:scope_0, :scope_1)" in sql + assert params["scope_0"] == 7 + assert params["scope_1"] == 9 diff --git a/tests/repository/test_search_reader.py b/tests/repository/test_search_reader.py index ed6da6e62..cccf2668b 100644 --- a/tests/repository/test_search_reader.py +++ b/tests/repository/test_search_reader.py @@ -233,7 +233,7 @@ async def test_built_in_adapter_reads_manifest_readiness_when_tracing(monkeypatc assert trace.readiness is readiness read_readiness.assert_awaited_once_with(session, SCOPE, "sqlite-vec", "fake:384") - adapter.search.assert_awaited_once_with([0.1], limit=5) + adapter.search.assert_awaited_once_with([0.1], limit=5, projects=SCOPE) @pytest.mark.asyncio diff --git a/tests/repository/test_search_trace.py b/tests/repository/test_search_trace.py index 89efb107d..055050960 100644 --- a/tests/repository/test_search_trace.py +++ b/tests/repository/test_search_trace.py @@ -68,11 +68,10 @@ def runtime_log_attrs(self) -> dict[str, Any]: class _TraceVectorIndex: - def __init__(self, project_id: int) -> None: + def __init__(self) -> None: self.matches: list[VectorMatch] = [] self.scope = VectorIndexScope( namespace="trace-test", - project_id=project_id, embedding_identity="trace-embedding", dimensions=4, ) @@ -80,16 +79,18 @@ def __init__(self, project_id: int) -> None: async def initialize(self) -> None: return None - async def upsert(self, records: Sequence[VectorRecord]) -> None: + async def upsert(self, project_id: int, records: Sequence[VectorRecord]) -> None: return None - async def delete(self, records: Sequence[VectorDeletion]) -> None: + async def delete(self, project_id: int, records: Sequence[VectorDeletion]) -> None: return None - async def delete_entity(self, entity_id: int) -> None: + async def delete_entity(self, project_id: int, entity_id: int) -> None: return None - async def search(self, query: Sequence[float], *, limit: int) -> list[VectorMatch]: + async def search( + self, query: Sequence[float], *, limit: int, projects: ProjectScope + ) -> list[VectorMatch]: return self.matches[:limit] @@ -525,7 +526,7 @@ def _repository( "semantic_vector_k": 10, } ) - vector_index = _TraceVectorIndex(test_project.id) + vector_index = _TraceVectorIndex() repository_type = ( PostgresSearchRepository if config.database_backend == DatabaseBackend.POSTGRES diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 4a5857aca..43fe5806f 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -148,7 +148,6 @@ class _RecordingVectorIndex: scope = VectorIndexScope( namespace="basic-memory-test", - project_id=1, embedding_identity="stub:4", dimensions=4, ) @@ -159,13 +158,13 @@ def __init__(self) -> None: async def initialize(self) -> None: return None - async def upsert(self, records: Sequence[VectorRecord]) -> None: + async def upsert(self, project_id: int, records: Sequence[VectorRecord]) -> None: self.upserted_records.extend(records) - async def delete(self, records: Sequence[VectorDeletion]) -> None: + async def delete(self, project_id: int, records: Sequence[VectorDeletion]) -> None: return None - async def delete_entity(self, entity_id: int) -> None: + async def delete_entity(self, project_id: int, entity_id: int) -> None: return None async def search( @@ -173,6 +172,7 @@ async def search( query: Sequence[float], *, limit: int, + projects: ProjectScope, ) -> list[VectorMatch]: return [] @@ -475,7 +475,9 @@ async def test_external_reconciliation_holds_project_lock_through_orphan_cleanup events: list[str] = [] adapter: Any = SimpleNamespace( scope=_RecordingVectorIndex.scope, - delete_orphans=AsyncMock(side_effect=lambda _live_keys: events.append("delete_orphans")), + delete_orphans=AsyncMock( + side_effect=lambda _project_id, _live_keys: events.append("delete_orphans") + ), ) repo._semantic_vector_index = adapter session = AsyncMock() @@ -508,7 +510,7 @@ async def fake_scoped_session(_session_maker): assert events == ["project_lock", "manifest_read", "delete_orphans", "commit"] adapter.delete_orphans.assert_awaited_once_with( - [VectorKey(entity_id=41, chunk_key="entity:41:0")] + 1, [VectorKey(entity_id=41, chunk_key="entity:41:0")] ) @@ -654,7 +656,9 @@ async def test_project_vector_cleanup_uses_available_adapter( events: list[str] = [] adapter: Any = SimpleNamespace( initialize=AsyncMock(side_effect=lambda: events.append("initialize")), - delete_entity=AsyncMock(side_effect=lambda _entity_id: events.append("delete")), + delete_entity=AsyncMock( + side_effect=lambda _project_id, _entity_id: events.append("delete") + ), ) repo._semantic_vector_index = adapter repo._semantic_vector_index_name = "milvus" @@ -697,8 +701,8 @@ async def fake_scoped_session(_session_maker): adapter.initialize.assert_awaited_once() assert adapter.delete_entity.await_args_list == [ - ((41,), {}), - ((42,), {}), + ((1, 41), {}), + ((1, 42), {}), ] expected_events = [ "project_lock", @@ -912,7 +916,9 @@ async def test_external_entity_cleanup_uses_matching_project_adapter(monkeypatch events: list[str] = [] adapter: Any = SimpleNamespace( initialize=AsyncMock(side_effect=lambda: events.append("initialize")), - delete_entity=AsyncMock(side_effect=lambda _entity_id: events.append("delete")), + delete_entity=AsyncMock( + side_effect=lambda _project_id, _entity_id: events.append("delete") + ), ) repo._semantic_vector_index = adapter repo._semantic_vector_index_name = "milvus" @@ -948,7 +954,7 @@ async def fake_scoped_session(_session_maker): ) adapter.initialize.assert_awaited_once() - assert adapter.delete_entity.await_args_list == [((41,), {}), ((42,), {})] + assert adapter.delete_entity.await_args_list == [((1, 41), {}), ((1, 42), {})] assert events == [ "project_lock", "ownership_read", diff --git a/tests/repository/test_semantic_vector_index.py b/tests/repository/test_semantic_vector_index.py index 96ec7b4f8..224ab0dd9 100644 --- a/tests/repository/test_semantic_vector_index.py +++ b/tests/repository/test_semantic_vector_index.py @@ -14,6 +14,7 @@ from basic_memory.repository.embedding_provider import EmbeddingProvider from basic_memory.repository.postgres_search_repository import PostgresSearchRepository from basic_memory.repository.search_repository import create_search_repository +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_errors import ( SemanticDependenciesMissingError, ) @@ -59,13 +60,13 @@ def __init__(self, scope: VectorIndexScope): async def initialize(self) -> None: return None - async def upsert(self, records: Sequence[VectorRecord]) -> None: + async def upsert(self, project_id: int, records: Sequence[VectorRecord]) -> None: return None - async def delete(self, records: Sequence[VectorDeletion]) -> None: + async def delete(self, project_id: int, records: Sequence[VectorDeletion]) -> None: return None - async def delete_entity(self, entity_id: int) -> None: + async def delete_entity(self, project_id: int, entity_id: int) -> None: return None async def search( @@ -73,6 +74,7 @@ async def search( query: Sequence[float], *, limit: int, + projects: ProjectScope, ) -> list[VectorMatch]: return [] @@ -92,7 +94,6 @@ def _postgres_config(**overrides: object) -> BasicMemoryConfig: def test_vector_contract_values_and_dimension_validation() -> None: scope = VectorIndexScope( namespace="basic-memory-test", - project_id=7, embedding_identity="stub:3", dimensions=3, ) @@ -127,9 +128,9 @@ def test_selector_defaults_to_pgvector_and_sqlite_remains_automatic() -> None: _postgres_config(semantic_vector_index="test-extension") -def test_scope_is_stable_credential_free_and_project_isolated() -> None: +def test_scope_is_stable_and_credential_free() -> None: provider: EmbeddingProvider = StubEmbeddingProvider() - first = build_vector_index_scope(_postgres_config(), provider, project_id=7) + first = build_vector_index_scope(_postgres_config(), provider) rotated_password = build_vector_index_scope( _postgres_config( database_url=( @@ -137,15 +138,12 @@ def test_scope_is_stable_credential_free_and_project_isolated() -> None: ) ), provider, - project_id=7, ) - other_project = build_vector_index_scope(_postgres_config(), provider, project_id=8) other_user = build_vector_index_scope( _postgres_config( database_url="postgresql+asyncpg://tenant-user:secret@db.example.test:5432/memory" ), provider, - project_id=7, ) other_schema = build_vector_index_scope( _postgres_config( @@ -155,7 +153,6 @@ def test_scope_is_stable_credential_free_and_project_isolated() -> None: ) ), provider, - project_id=7, ) first_socket = build_vector_index_scope( _postgres_config( @@ -164,7 +161,6 @@ def test_scope_is_stable_credential_free_and_project_isolated() -> None: ) ), provider, - project_id=7, ) other_socket = build_vector_index_scope( _postgres_config( @@ -173,18 +169,16 @@ def test_scope_is_stable_credential_free_and_project_isolated() -> None: ) ), provider, - project_id=7, ) assert first.namespace == rotated_password.namespace assert "secret" not in first.namespace - assert first.project_id != other_project.project_id assert first.namespace != other_user.namespace assert first.namespace != other_schema.namespace assert first_socket.namespace != other_socket.namespace assert first.embedding_identity == "StubEmbeddingProvider:stub-model:3" assert first.dimensions == 3 - assert first.storage_key == rotated_password.storage_key + assert first == rotated_password def test_milvus_without_optional_dependencies_reports_install_extra(monkeypatch) -> None: @@ -213,7 +207,6 @@ def import_without_pymilvus( ): create_semantic_vector_index( session_maker=MagicMock(), - project_id=7, app_config=config, database_backend=DatabaseBackend.POSTGRES, embedding_provider=StubEmbeddingProvider(), @@ -222,7 +215,7 @@ def import_without_pymilvus( def test_search_repository_composition_root_injects_selected_adapter(monkeypatch) -> None: provider = StubEmbeddingProvider() - scope = build_vector_index_scope(_postgres_config(), provider, project_id=7) + scope = build_vector_index_scope(_postgres_config(), provider) index = StubVectorIndex(scope) monkeypatch.setattr( "basic_memory.repository.search_repository.create_embedding_provider", diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index f981fac1c..c53291b26 100644 --- a/tests/repository/test_sqlite_vector_search_repository.py +++ b/tests/repository/test_sqlite_vector_search_repository.py @@ -18,6 +18,7 @@ from basic_memory.repository.prefixing_provider import PrefixingEmbeddingProvider from basic_memory.repository import search_repository_base as search_repository_base_module from basic_memory.repository.search_index_row import SearchIndexRow +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_errors import SemanticVectorIndexExtensionError from basic_memory.repository.semantic_vector_index import ( VectorDeletion, @@ -76,7 +77,6 @@ class RecordingVectorIndex: def __init__(self) -> None: self.scope = VectorIndexScope( namespace="test", - project_id=1, embedding_identity="test", dimensions=4, ) @@ -91,20 +91,20 @@ def __init__(self) -> None: async def initialize(self) -> None: return None - async def upsert(self, records: Sequence[VectorRecord]) -> None: + async def upsert(self, project_id: int, records: Sequence[VectorRecord]) -> None: self.upsert_calls.append(list(records)) if self.fail_upsert: raise RuntimeError("adapter write failed") self.records.update({record.key: record.values for record in records}) - async def delete(self, records: Sequence[VectorDeletion]) -> None: + async def delete(self, project_id: int, records: Sequence[VectorDeletion]) -> None: self.deleted_entities.extend(sorted({record.key.entity_id for record in records})) if self.fail_delete_entity: raise RuntimeError("adapter delete failed") for record in records: self.records.pop(record.key, None) - async def delete_entity(self, entity_id: int) -> None: + async def delete_entity(self, project_id: int, entity_id: int) -> None: self.deleted_entities.append(entity_id) if self.fail_delete_entity: raise RuntimeError("adapter delete failed") @@ -112,7 +112,7 @@ async def delete_entity(self, entity_id: int) -> None: key: values for key, values in self.records.items() if key.entity_id != entity_id } - async def delete_orphans(self, live_keys: Sequence[VectorKey]) -> None: + async def delete_orphans(self, project_id: int, live_keys: Sequence[VectorKey]) -> None: self.reconcile_calls.append(list(live_keys)) live_key_set = set(live_keys) self.records = {key: values for key, values in self.records.items() if key in live_key_set} @@ -122,6 +122,7 @@ async def search( query: Sequence[float], *, limit: int, + projects: ProjectScope, ) -> list[VectorMatch]: if self.fail_search: raise RuntimeError("adapter query failed") @@ -436,7 +437,7 @@ async def test_sqlite_vec_reconciliation_is_project_scoped(search_repository): ) await session.commit() - await index.delete_orphans([]) + await index.delete_orphans(search_repository.project_id, []) async with db.scoped_session(search_repository.session_maker) as session: remaining = await session.execute( @@ -448,6 +449,71 @@ async def test_sqlite_vec_reconciliation_is_project_scoped(search_repository): assert remaining.scalars().all() == [902, 903] +@pytest.mark.asyncio +async def test_sqlite_vec_search_reads_every_project_in_scope(search_repository): + """One statement answers a multi-project scope; a single-project scope stays isolated.""" + if not isinstance(search_repository, SQLiteSearchRepository): + pytest.skip("sqlite-vec search behavior is local SQLite-only.") + + _enable_semantic(search_repository) + await search_repository.init_search_index() + index = cast(SQLiteVecIndex, search_repository._semantic_vector_index) + embedding_identity = search_repository._embedding_model_key() + own_project = search_repository.project_id + other_project = own_project + 1 + + async with db.scoped_session(search_repository.session_maker) as session: + await index._ensure_loaded(session) + await session.execute( + text( + "INSERT INTO search_vector_chunks (" + "id, entity_id, project_id, chunk_key, chunk_text, source_hash, " + "entity_fingerprint, embedding_model, vector_index, embedding_status" + ") VALUES (" + ":id, :entity_id, :project_id, :chunk_key, 'text', 'hash', " + "'fingerprint', :embedding_model, 'sqlite-vec', 'ready')" + ), + [ + { + "id": 911, + "entity_id": 911, + "project_id": own_project, + "chunk_key": "entity:911:0", + "embedding_model": embedding_identity, + }, + { + "id": 912, + "entity_id": 912, + "project_id": other_project, + "chunk_key": "entity:912:0", + "embedding_model": embedding_identity, + }, + ], + ) + await session.execute( + text( + "INSERT INTO search_vector_embeddings (rowid, embedding, source_hash) " + "VALUES (:rowid, :embedding, 'hash')" + ), + [ + {"rowid": 911, "embedding": "[1,0,0,0]"}, + {"rowid": 912, "embedding": "[0,1,0,0]"}, + ], + ) + await session.commit() + + query = [1.0, 0.0, 0.0, 0.0] + both = await index.search( + query, limit=10, projects=ProjectScope.of([other_project, own_project]) + ) + own_only = await index.search(query, limit=10, projects=ProjectScope.single(own_project)) + nothing = await index.search(query, limit=10, projects=ProjectScope.of([])) + + assert [match.key.entity_id for match in both] == [911, 912] + assert [match.key.entity_id for match in own_only] == [911] + assert nothing == [] + + @pytest.mark.asyncio async def test_sqlite_vec_delete_requires_pending_source_generation(search_repository): """A stale delete cannot remove a same-source vector that is already ready.""" @@ -486,7 +552,7 @@ async def test_sqlite_vec_delete_requires_pending_source_generation(search_repos await session.commit() deletion = VectorDeletion(key=key, source_hash="hash") - await index.delete([deletion]) + await index.delete(search_repository.project_id, [deletion]) async with db.scoped_session(search_repository.session_maker) as session: assert ( await session.scalar( @@ -499,7 +565,7 @@ async def test_sqlite_vec_delete_requires_pending_source_generation(search_repos ) await session.commit() - await index.delete([deletion]) + await index.delete(search_repository.project_id, [deletion]) async with db.scoped_session(search_repository.session_maker) as session: vector_count = await session.scalar( text("SELECT COUNT(*) FROM search_vector_embeddings WHERE rowid = 907") @@ -1459,19 +1525,19 @@ async def fake_scoped_session(_session_maker): monkeypatch.setattr(index, "_ensure_loaded", AsyncMock()) query_embedding = [0.1] * search_repository._vector_dimensions - await index.search(query_embedding, limit=10000) + await index.search(query_embedding, limit=10000, projects=search_repository.scope) assert captured_params == [ { "query": "[0.1, 0.1, 0.1, 0.1]", "vector_k": SQLITE_VEC_MAX_K, - "project_id": search_repository.project_id, + "scope_0": search_repository.project_id, "embedding_identity": search_repository._embedding_model_key(), "limit": 10000, } ] captured_params.clear() - await index.search(query_embedding, limit=500) + await index.search(query_embedding, limit=500, projects=search_repository.scope) assert captured_params[0]["vector_k"] == 500 assert captured_params[0]["limit"] == 500 diff --git a/tests/repository/test_vector_manifest_generation_ownership.py b/tests/repository/test_vector_manifest_generation_ownership.py index f92abbd0d..62f502476 100644 --- a/tests/repository/test_vector_manifest_generation_ownership.py +++ b/tests/repository/test_vector_manifest_generation_ownership.py @@ -13,6 +13,7 @@ from basic_memory.config import BasicMemoryConfig, DatabaseBackend from basic_memory.repository.postgres_search_repository import PostgresSearchRepository from basic_memory.repository.search_index_row import SearchIndexRow +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_vector_index import ( VectorDeletion, VectorIndexScope, @@ -74,17 +75,17 @@ def scope(self) -> VectorIndexScope: async def initialize(self) -> None: return None - async def upsert(self, records: Sequence[VectorRecord]) -> None: + async def upsert(self, project_id: int, records: Sequence[VectorRecord]) -> None: for record in records: self.records[record.key] = record - async def delete(self, records: Sequence[VectorDeletion]) -> None: + async def delete(self, project_id: int, records: Sequence[VectorDeletion]) -> None: for deletion in records: current = self.records.get(deletion.key) if current is not None and current.source_hash == deletion.source_hash: self.records.pop(deletion.key) - async def delete_entity(self, entity_id: int) -> None: + async def delete_entity(self, project_id: int, entity_id: int) -> None: self.records = { key: record for key, record in self.records.items() if key.entity_id != entity_id } @@ -94,6 +95,7 @@ async def search( query: Sequence[float], *, limit: int, + projects: ProjectScope, ) -> list[VectorMatch]: return [] @@ -144,7 +146,6 @@ async def _repositories( vector_index = InMemoryExternalVectorIndex( VectorIndexScope( namespace="generation-ownership", - project_id=project_id, embedding_identity="test:4", dimensions=4, ) From 06b5f935f796f765ac48488a7452e23079765533 Mon Sep 17 00:00:00 2001 From: Paul Hernandez <60959+phernandez@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:55:34 -0500 Subject: [PATCH 5/8] feat(api): search an explicit set of projects with one query (#1567) Signed-off-by: phernandez Co-authored-by: Claude Fable 5.1 --- CHANGELOG.md | 10 + src/basic_memory/api/app.py | 3 + src/basic_memory/api/v2/routers/__init__.py | 2 + .../api/v2/routers/scoped_search_router.py | 114 +++ .../api/v2/routers/search_router.py | 37 +- src/basic_memory/api/v2/utils.py | 44 +- src/basic_memory/repository/search_reader.py | 5 + .../repository/search_repository.py | 63 ++ .../repository/sqlite_vec_index.py | 13 +- src/basic_memory/schemas/search.py | 19 +- .../services/scoped_search_service.py | 156 +++++ src/basic_memory/services/search_service.py | 241 ++++--- tests/api/v2/test_scoped_search_router.py | 657 ++++++++++++++++++ tests/api/v2/test_search_router_telemetry.py | 5 +- tests/repository/test_search_reader.py | 17 + .../repository/test_search_reader_factory.py | 115 +++ tests/repository/test_sqlite_vec_index.py | 41 ++ 17 files changed, 1408 insertions(+), 134 deletions(-) create mode 100644 src/basic_memory/api/v2/routers/scoped_search_router.py create mode 100644 src/basic_memory/services/scoped_search_service.py create mode 100644 tests/api/v2/test_scoped_search_router.py create mode 100644 tests/repository/test_search_reader_factory.py create mode 100644 tests/repository/test_sqlite_vec_index.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 0d347cd53..a53634749 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,16 @@ ### Features +- **#1558**: `QUERY /v2/search/` (and `POST /v2/search/` for clients that cannot send + QUERY) searches an explicit set of projects in one database with one query. The body + is the project search body plus `project_ids`, a required list of internal ids the + caller has already authorized; an empty list answers no rows and there is no way to + ask for every project. Full-text, vector, and hybrid retrieval run the same reader + the project route runs, so one project here ranks exactly as its own route does, and + results are one ranking over the union rather than merged per-project pages. Hits are + hydrated only from projects in scope. Every search result, on both routes, now + carries `project_id` and `project_external_id`. + - **#1512**: Word, PowerPoint, and CSV files get the same sidecar Markdown note a PDF gets. `bm import document ` indexes the project, extracts the file, and writes `..md` next to it plus a run note under diff --git a/src/basic_memory/api/app.py b/src/basic_memory/api/app.py index 5e2cb8cf6..c2779448b 100644 --- a/src/basic_memory/api/app.py +++ b/src/basic_memory/api/app.py @@ -14,6 +14,7 @@ knowledge_router as v2_knowledge, project_router as v2_project, memory_router as v2_memory, + scoped_search_router as v2_scoped_search, search_router as v2_search, resource_router as v2_resource, directory_router as v2_directory, @@ -136,6 +137,8 @@ async def workspace_permalink_context_middleware(request: Request, call_next): app.include_router(v2_schema, prefix="/v2/projects/{project_id}") app.include_router(v2_inspect, prefix="/v2/projects/{project_id}") app.include_router(v2_project, prefix="/v2") +# Database-scoped search: one query over an explicit set of projects. +app.include_router(v2_scoped_search, prefix="/v2") # Legacy web app proxy paths (compat with /proxy/projects/projects) app.include_router(v2_project, prefix="/proxy/projects") diff --git a/src/basic_memory/api/v2/routers/__init__.py b/src/basic_memory/api/v2/routers/__init__.py index 29015ba37..9c0bc36ab 100644 --- a/src/basic_memory/api/v2/routers/__init__.py +++ b/src/basic_memory/api/v2/routers/__init__.py @@ -6,6 +6,7 @@ from basic_memory.api.v2.routers.knowledge_router import router as knowledge_router from basic_memory.api.v2.routers.project_router import router as project_router from basic_memory.api.v2.routers.memory_router import router as memory_router +from basic_memory.api.v2.routers.scoped_search_router import router as scoped_search_router from basic_memory.api.v2.routers.search_router import router as search_router from basic_memory.api.v2.routers.resource_router import router as resource_router from basic_memory.api.v2.routers.directory_router import router as directory_router @@ -22,6 +23,7 @@ "knowledge_router", "project_router", "memory_router", + "scoped_search_router", "search_router", "resource_router", "directory_router", diff --git a/src/basic_memory/api/v2/routers/scoped_search_router.py b/src/basic_memory/api/v2/routers/scoped_search_router.py new file mode 100644 index 000000000..a6fe02a54 --- /dev/null +++ b/src/basic_memory/api/v2/routers/scoped_search_router.py @@ -0,0 +1,114 @@ +"""Database-scoped search: one query over an explicit set of projects. + +The caller names the projects. Nothing here discovers a scope, and an empty set +answers no rows; Cloud resolves authorization first and passes the effective ids. +The route shares the project route's reader, so one project here ranks exactly as +that project's own route does, and several projects are one ranking over the union +rather than per-project pages merged afterwards. +""" + +import asyncio + +import logfire +from fastapi import APIRouter, Query, Response + +from basic_memory import db +from basic_memory.api.v2.utils import ( + load_temporal_metadata, + search_error_boundary, + to_search_results, +) +from basic_memory.deps import AppConfigDep, SessionMakerDep +from basic_memory.repository.search_repository import create_search_reader +from basic_memory.repository.search_scope import ProjectScope +from basic_memory.schemas.search import ScopedSearchQuery, SearchResponse, SearchRetrievalMode +from basic_memory.services.scoped_search_service import ScopedSearchService + +# App registration mounts this router at /v2. +router = APIRouter(tags=["search"]) + + +@router.api_route( + "/search/", + methods=["QUERY"], + response_model=SearchResponse, + include_in_schema=False, +) +@router.post("/search/", response_model=SearchResponse) +async def search_scope( + query: ScopedSearchQuery, + app_config: AppConfigDep, + session_maker: SessionMakerDep, + response: Response, + page: int = Query(1, ge=1), + page_size: int = Query(10, ge=1, le=1000), +) -> SearchResponse: + """Search an explicit set of projects in this database. + + Results are one ranking across the whole scope and carry ``project_id`` and + ``project_external_id``. Hydration reads only from projects in scope. + """ + response.headers["Accept-Query"] = "application/json" + # The read cache is keyed by one project's generation. A set of projects has no + # single generation to invalidate on, so this route is never cached. + response.headers["Cache-Control"] = "no-store" + + scope = ProjectScope.of(query.project_ids) + service = ScopedSearchService( + session_maker, scope, create_search_reader(session_maker, scope, app_config) + ) + temporal_requested = query.has_temporal_filter() + exact_count_available = query.retrieval_mode == SearchRetrievalMode.FTS + offset = (page - 1) * page_size + + with logfire.span( + "api.request.search_scope", + entrypoint="api", + domain="search", + action="search_scope", + project_count=len(scope.project_ids), + page=page, + page_size=page_size, + retrieval_mode=query.retrieval_mode.value, + has_temporal_filter=temporal_requested, + ): + with search_error_boundary(): + if exact_count_available: + results, total = await asyncio.gather( + service.search(query, limit=page_size, offset=offset), + service.count(query), + ) + has_more = offset + len(results) < total + else: + # Trigger: semantic modes would need another vector or hybrid pass to count. + # Why: a search should not pay for a second semantic retrieval. + # Outcome: probe one row past the page, leave total at 0, mark it inexact. + results = await service.search(query, limit=page_size + 1, offset=offset) + total = 0 + has_more = len(results) > page_size + results = results[:page_size] + + temporal_by_source = {} + project_external_ids: dict[int, str] = {} + if results: + async with db.scoped_session(session_maker) as session: + project_external_ids = await service.project_external_ids(session, results) + if temporal_requested: + temporal_by_source = await load_temporal_metadata(service, session, results) + search_results = await to_search_results( + service, + results, + temporal_by_source=temporal_by_source, + project_external_ids=project_external_ids, + ) + return SearchResponse( + results=search_results, + current_page=page, + page_size=page_size, + total=total, + total_is_exact=exact_count_available, + has_more=has_more, + # None, not False, when nothing was asked: an ordinary search payload stays + # exactly what it was before valid time existed. + temporal_applied=True if temporal_requested else None, + ) diff --git a/src/basic_memory/api/v2/routers/search_router.py b/src/basic_memory/api/v2/routers/search_router.py index 2e8460142..f5c137486 100644 --- a/src/basic_memory/api/v2/routers/search_router.py +++ b/src/basic_memory/api/v2/routers/search_router.py @@ -9,11 +9,15 @@ from contextlib import nullcontext from typing import Annotated -from fastapi import APIRouter, Depends, HTTPException, Path, Response +from fastapi import APIRouter, Depends, Path, Response import logfire from basic_memory import db -from basic_memory.api.v2.utils import load_temporal_metadata, to_search_results +from basic_memory.api.v2.utils import ( + load_temporal_metadata, + search_error_boundary, + to_search_results, +) from basic_memory.deps import ( EntityServiceV2ExternalDep, MemoryTimeIndexRepositoryV2ExternalDep, @@ -32,12 +36,6 @@ read_cache_request_digest, ) from basic_memory.read_cache.policy import SEARCH_READ_CACHE_TTL_SECONDS -from basic_memory.repository.semantic_errors import ( - RerankProviderContractError, - RerankTransientError, - SemanticDependenciesMissingError, - SemanticSearchDisabledError, -) from basic_memory.schemas.search import SearchQuery, SearchResponse, SearchRetrievalMode from basic_memory.services.search_guidance import unspaced_script_query_hint @@ -91,6 +89,7 @@ async def search( session_maker: SessionMakerDep, read_cache: SearchReadCacheDep, response: Response, + internal_project_id: ProjectExternalIdPathDep, project_id: str = Path(..., description="Project external UUID"), page: int = 1, page_size: int = 10, @@ -157,7 +156,7 @@ async def search( offset = (page - 1) * page_size exact_count_available = query.retrieval_mode == SearchRetrievalMode.FTS - try: + with search_error_boundary(): with logfire.span( "api.search.search.execute_query", domain="search", @@ -176,21 +175,6 @@ async def search( query, limit=page_size + 1, offset=offset ) total = 0 - except SemanticSearchDisabledError as exc: - raise HTTPException(status_code=400, detail=str(exc)) from exc - except SemanticDependenciesMissingError as exc: - raise HTTPException(status_code=400, detail=str(exc)) from exc - except RerankTransientError as exc: - # Returning raw retrieval order would make pagination inconsistent with - # earlier reranked pages. Preserve ordering semantics and make the outage - # explicitly retryable instead. - raise HTTPException(status_code=503, detail=str(exc)) from exc - except RerankProviderContractError as exc: - # Upstream reranker returned a malformed response — an upstream fault, not a - # client error and not a transient outage (those map to a retryable 503). - raise HTTPException(status_code=502, detail=str(exc)) from exc - except ValueError as exc: - raise HTTPException(status_code=400, detail=str(exc)) from exc with logfire.span( "api.search.search.paginate_results", @@ -228,7 +212,10 @@ async def search( temporal_repository, session, results ) search_results = await to_search_results( - entity_service, results, temporal_by_source=temporal_by_source + entity_service, + results, + temporal_by_source=temporal_by_source, + project_external_ids={internal_project_id: project_id}, ) with logfire.span( "api.search.search.build_response", diff --git a/src/basic_memory/api/v2/utils.py b/src/basic_memory/api/v2/utils.py index c9d0350bd..82c7942f9 100644 --- a/src/basic_memory/api/v2/utils.py +++ b/src/basic_memory/api/v2/utils.py @@ -1,11 +1,19 @@ from collections import defaultdict -from collections.abc import Mapping +from collections.abc import Iterator, Mapping +from contextlib import contextmanager from typing import Any, Protocol, Optional, List, Sequence import logfire +from fastapi import HTTPException from sqlalchemy.ext.asyncio import AsyncSession from basic_memory.models import MemoryTimeIndex from basic_memory.repository.search_repository import SearchIndexRow +from basic_memory.repository.semantic_errors import ( + RerankProviderContractError, + RerankTransientError, + SemanticDependenciesMissingError, + SemanticSearchDisabledError, +) from basic_memory.schemas.memory import ( EntitySummary, ObservationSummary, @@ -54,6 +62,28 @@ async def find_for_sources( type TemporalMetadataBySource = Mapping[tuple[str, int], list[TemporalResultMetadata]] +@contextmanager +def search_error_boundary() -> Iterator[None]: + """Map search failures onto HTTP statuses the same way on every search route.""" + try: + yield + except SemanticSearchDisabledError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + except SemanticDependenciesMissingError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + except RerankTransientError as exc: + # Returning raw retrieval order would make pagination inconsistent with + # earlier reranked pages. Preserve ordering semantics and make the outage + # explicitly retryable instead. + raise HTTPException(status_code=503, detail=str(exc)) from exc + except RerankProviderContractError as exc: + # Upstream reranker returned a malformed response: an upstream fault, not a + # client error and not a transient outage (those map to a retryable 503). + raise HTTPException(status_code=502, detail=str(exc)) from exc + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + async def get_entities_by_id_lookup( entity_service: EntityServiceBatchLookup, entity_ids: Sequence[int], @@ -304,7 +334,13 @@ async def to_search_results( results: List[SearchIndexRow], *, temporal_by_source: TemporalMetadataBySource | None = None, + project_external_ids: Mapping[int, str] | None = None, ) -> list[SearchResult]: + """Shape one page of search rows into the public result. + + ``project_external_ids`` maps each row's project to its external id; a caller + that knows the projects passes it and results carry both identities. + """ with logfire.span( "search.hydrate_results", domain="search", @@ -391,6 +427,12 @@ async def to_search_results( if temporal_by_source else None ), + project_id=result.project_id, + project_external_id=( + project_external_ids.get(result.project_id) + if project_external_ids + else None + ), ) ) return search_results diff --git a/src/basic_memory/repository/search_reader.py b/src/basic_memory/repository/search_reader.py index 0b5c81eb5..8d5d80b28 100644 --- a/src/basic_memory/repository/search_reader.py +++ b/src/basic_memory/repository/search_reader.py @@ -615,6 +615,9 @@ async def vector_only( ``candidate_limit`` is supplied only by a composed retrieval stage that already sized the shared candidate pool. """ + # An empty scope admits no rows; embedding the query would buy nothing. + if self.scope.is_empty: + return [] query_text = (query.search_text or "").strip() if candidate_limit is None: candidate_limit = self._candidate_limit(limit, offset, query_text) @@ -833,6 +836,8 @@ async def hybrid( ``max(vec, fts) + FUSION_BONUS * min(vec, fts)`` preserves the dominant signal and rewards dual-source agreement. """ + if self.scope.is_empty: + return [] query_text = (query.search_text or "").strip() rerank = self._active_rerank(query_text) query_start = time.perf_counter() diff --git a/src/basic_memory/repository/search_repository.py b/src/basic_memory/repository/search_repository.py index b21aca5b1..ca4eb0c14 100644 --- a/src/basic_memory/repository/search_repository.py +++ b/src/basic_memory/repository/search_repository.py @@ -16,14 +16,25 @@ from basic_memory.config import BasicMemoryConfig, DatabaseBackend from basic_memory.repository.embedding_provider_factory import create_embedding_provider from basic_memory.repository.rerank_provider_factory import create_rerank_provider +from basic_memory.repository.postgres_search_query import PostgresFts from basic_memory.repository.postgres_search_repository import PostgresSearchRepository +from basic_memory.repository.search_filters import FtsBackend from basic_memory.repository.search_index_row import SearchIndexRow +from basic_memory.repository.search_reader import ( + Reranking, + SearchReader, + SemanticSearch, + VectorRetrieval, +) from basic_memory.repository.search_repository_base import ChunkManifestRow, SearchIndexKey +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.search_trace import SearchTraceCollector from basic_memory.repository.semantic_vector_index_factory import ( create_semantic_vector_index, resolve_semantic_vector_index_name, + semantic_embedding_identity, ) +from basic_memory.repository.sqlite_search_query import SQLiteFts from basic_memory.runtime.vector_sync import VectorSyncBatchResult from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode @@ -265,8 +276,60 @@ def create_search_repository( ) +def create_search_reader( + session_maker: async_sessionmaker[AsyncSession], + scope: ProjectScope, + app_config: BasicMemoryConfig, + database_backend: Optional[DatabaseBackend] = None, +) -> SearchReader: + """Compose the read path over an explicit scope, with no project repository. + + Resolves the same shared embedding provider, vector adapter, and reranker that + ``create_search_repository`` hands a project repository, so a scoped search runs + the pipeline a project's own route runs, over a wider scope. Whether semantic + retrieval is available is decided here, once, from configuration; a semantic + query against a reader built without it fails as a disabled feature. + """ + backend = database_backend or app_config.database_backend + fts: FtsBackend = ( + PostgresFts(session_maker) + if backend == DatabaseBackend.POSTGRES + else SQLiteFts(session_maker) + ) + if not app_config.semantic_search_enabled: + return SearchReader(scope, fts) + + embedding_provider = create_embedding_provider(app_config) + vector_index_name, vector_index = create_semantic_vector_index( + session_maker=session_maker, + app_config=app_config, + database_backend=backend, + embedding_provider=embedding_provider, + ) + vector = VectorRetrieval( + index=vector_index, + index_name=vector_index_name, + embedding_provider=embedding_provider, + embedding_model=semantic_embedding_identity(embedding_provider), + vector_k=app_config.semantic_vector_k, + min_similarity=app_config.semantic_min_similarity, + ) + rerank_provider = create_rerank_provider(app_config) + rerank = ( + Reranking( + provider=rerank_provider, + candidates=app_config.reranker_candidates, + max_document_chars=app_config.reranker_max_document_chars, + ) + if rerank_provider is not None + else None + ) + return SearchReader(scope, fts, SemanticSearch(session_maker, scope, fts, vector, rerank)) + + __all__ = [ "SearchRepository", "SearchIndexRow", + "create_search_reader", "create_search_repository", ] diff --git a/src/basic_memory/repository/sqlite_vec_index.py b/src/basic_memory/repository/sqlite_vec_index.py index f85f40248..cfa45ddf7 100644 --- a/src/basic_memory/repository/sqlite_vec_index.py +++ b/src/basic_memory/repository/sqlite_vec_index.py @@ -79,7 +79,18 @@ async def _ensure_loaded(self, session: AsyncSession) -> None: "basic-memory under uv-managed or Homebrew Python, or disable " "semantic search." ) - await driver_connection.enable_load_extension(True) + try: + await driver_connection.enable_load_extension(True) + except AttributeError as exc: + # aiosqlite exposes the wrapper method even when the wrapped + # sqlite3.Connection was built without extension support, so + # calling it is the authoritative probe (#711). + raise SemanticDependenciesMissingError( + "This Python build does not support SQLite extension loading " + "(no enable_load_extension on sqlite3.Connection). Reinstall " + "basic-memory under uv-managed or Homebrew Python, or disable " + "semantic search." + ) from exc await driver_connection.load_extension(sqlite_vec.loadable_path()) await driver_connection.enable_load_extension(False) await session.execute(text("SELECT vec_version()")) diff --git a/src/basic_memory/schemas/search.py b/src/basic_memory/schemas/search.py index ba2ce0cc1..ddd1435a0 100644 --- a/src/basic_memory/schemas/search.py +++ b/src/basic_memory/schemas/search.py @@ -6,7 +6,7 @@ 3. Full-text search across content """ -from typing import Optional, List, Union, Any +from typing import Annotated, Optional, List, Union, Any from datetime import datetime from enum import Enum from pydantic import BaseModel, Field, ValidationInfo, field_validator, model_validator @@ -227,6 +227,17 @@ def has_boolean_operators(self) -> bool: return any(pattern in text for pattern in boolean_patterns) +class ScopedSearchQuery(SearchQuery): + """A search over an explicit set of projects in one database. + + ``project_ids`` are internal ids the caller has already authorized; the caller + decides what is visible and this route never widens it. The set is required and + may be empty, which answers no rows. There is no spelling for every project. + """ + + project_ids: list[Annotated[int, Field(strict=True, gt=0)]] + + class TemporalRangeValue(BaseModel): """One authored interval, as a caller sees it. @@ -299,6 +310,12 @@ class SearchResult(BaseModel): # of multiple kinds must not be a schema break later. temporal: Optional[List[TemporalResultMetadata]] = None + # The project this hit belongs to. The project route fills both from its path; + # the database-scoped route fills them from each row, since one page can span + # several projects. + project_id: Optional[int] = None + project_external_id: Optional[str] = None + class SearchResponse(BaseModel): """Wrapper for search results.""" diff --git a/src/basic_memory/services/scoped_search_service.py b/src/basic_memory/services/scoped_search_service.py new file mode 100644 index 000000000..38900b618 --- /dev/null +++ b/src/basic_memory/services/scoped_search_service.py @@ -0,0 +1,156 @@ +"""Search over an explicit set of projects in one database. + +The project route and this service run the same ``SearchReader``; only the scope +differs. Hydration stays inside the scope too: owning entities, relation endpoints, +project identities, and valid-time assertions are read only from the projects the +search was allowed to read, so a page never names something outside its scope. +""" + +from collections.abc import Iterable, Sequence + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker +from sqlalchemy.orm import load_only + +from basic_memory import db +from basic_memory.models import Entity, MemoryTimeIndex, Project +from basic_memory.repository.repository import SELECT_BY_IDS_CHUNK_SIZE +from basic_memory.repository.search_index_row import SearchIndexRow +from basic_memory.repository.search_reader import SearchReader +from basic_memory.repository.search_scope import ProjectScope +from basic_memory.schemas.search import SearchQuery +from basic_memory.services.search_service import ( + include_legacy_note_type_spellings, + prepare_search_query, + relaxed_fts_fallback_eligible, +) + + +def _chunks[T](values: Sequence[T]) -> Iterable[Sequence[T]]: + """Split a bind list at the shared per-statement parameter bound.""" + for start in range(0, len(values), SELECT_BY_IDS_CHUNK_SIZE): + yield values[start : start + SELECT_BY_IDS_CHUNK_SIZE] + + +class ScopedSearchService: + """Run and hydrate searches over one ``ProjectScope``.""" + + def __init__( + self, + session_maker: async_sessionmaker[AsyncSession], + scope: ProjectScope, + reader: SearchReader, + ) -> None: + self.session_maker = session_maker + self.scope = scope + self.reader = reader + + # --- Retrieval --- + + async def search( + self, + query: SearchQuery, + *, + limit: int, + offset: int, + ) -> list[SearchIndexRow]: + """One ranking over every project in scope.""" + prepared = prepare_search_query(query) + if prepared is None: + return [] + prepared = await include_legacy_note_type_spellings( + self.session_maker, self.scope, prepared + ) + allow_relaxed = relaxed_fts_fallback_eligible( + query, prepared.search_text, prepared.retrieval_mode + ) + return await self.reader.search( + prepared, limit=limit, offset=offset, allow_relaxed=allow_relaxed + ) + + async def count(self, query: SearchQuery) -> int: + """Exact full-text match count over every project in scope.""" + prepared = prepare_search_query(query) + if prepared is None: + return 0 + prepared = await include_legacy_note_type_spellings( + self.session_maker, self.scope, prepared + ) + allow_relaxed = relaxed_fts_fallback_eligible( + query, prepared.search_text, prepared.retrieval_mode + ) + return await self.reader.count(prepared, allow_relaxed=allow_relaxed) + + # --- Hydration, bounded by the scope --- + + async def get_entities_by_id(self, ids: Sequence[int]) -> Sequence[Entity]: + """The entities a page of hits refers to, read only from projects in scope. + + Only the fields result shaping reads are loaded. Ids reached through a scoped + search already belong to the scope; the predicate keeps that true by + construction rather than by trust. + """ + if not ids or self.scope.is_empty: + return [] + entities: list[Entity] = [] + async with db.scoped_session(self.session_maker) as session: + for chunk in _chunks(list(ids)): + result = await session.scalars( + select(Entity) + .where( + Entity.project_id.in_(self.scope.project_ids), + Entity.id.in_(chunk), + ) + .options( + load_only( + Entity.id, Entity.project_id, Entity.permalink, Entity.external_id + ) + ) + ) + entities.extend(result.all()) + return entities + + async def find_for_sources( + self, + session: AsyncSession, + sources: Iterable[tuple[str, int]], + ) -> Sequence[MemoryTimeIndex]: + """The valid-time assertions behind a page of hits, read only from projects in scope. + + Mirrors ``MemoryTimeIndexRepository.find_for_sources`` for a set of projects: + one statement per source type, chunked at the bind bound. + """ + if self.scope.is_empty: + return [] + ids_by_type: dict[str, list[int]] = {} + for source_type, source_id in sources: + ids_by_type.setdefault(source_type, []).append(source_id) + + rows: list[MemoryTimeIndex] = [] + for source_type, source_ids in ids_by_type.items(): + for chunk in _chunks(source_ids): + result = await session.scalars( + select(MemoryTimeIndex) + .where( + MemoryTimeIndex.project_id.in_(self.scope.project_ids), + MemoryTimeIndex.source_type == source_type, + MemoryTimeIndex.source_id.in_(chunk), + ) + .order_by(MemoryTimeIndex.source_id, MemoryTimeIndex.id) + ) + rows.extend(result.all()) + return rows + + async def project_external_ids( + self, + session: AsyncSession, + rows: Sequence[SearchIndexRow], + ) -> dict[int, str]: + """External ids for the projects a page of hits came from.""" + project_ids = sorted({row.project_id for row in rows}) + if not project_ids: + return {} + result = await session.execute( + select(Project.id, Project.external_id).where(Project.id.in_(project_ids)) + ) + return {int(project_id): str(external_id) for project_id, external_id in result.all()} diff --git a/src/basic_memory/services/search_service.py b/src/basic_memory/services/search_service.py index 0a0f70317..1ba056858 100644 --- a/src/basic_memory/services/search_service.py +++ b/src/basic_memory/services/search_service.py @@ -11,6 +11,7 @@ from dateparser import parse from fastapi import BackgroundTasks from loguru import logger +from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker import logfire @@ -24,6 +25,7 @@ SearchRepository, ) from basic_memory.repository.search_query import PreparedSearchQuery, relaxed_query_words +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.search_trace import SearchTraceCollector from basic_memory.schemas.base import normalize_note_type from basic_memory.schemas.search import SearchQuery, SearchItemType, SearchRetrievalMode @@ -140,6 +142,132 @@ def _strip_nul(value: str) -> str: return value.replace("\x00", "") +def prepare_search_query(query: SearchQuery) -> PreparedSearchQuery | None: + """Normalize a ``SearchQuery`` into the prepared form every reader consumes. + + Returns ``None`` when the query names no criteria at all, so callers can answer + an empty page without touching storage. + """ + search_text = query.text + tags = query.tags + + # Support tag: shorthand by mapping to tags filter. + if search_text is not None: + search_text = search_text.strip() or None + if search_text and search_text.lower().startswith("tag:"): + tag_values = re.split(r"[,\s]+", search_text[4:].strip()) + parsed_tags = [t for t in tag_values if t] + if parsed_tags: + tags = parsed_tags + search_text = None + + after_date = ( + (query.after_date if isinstance(query.after_date, datetime) else parse(query.after_date)) + if query.after_date + else None + ) + + # Merge structured metadata filters (explicit + convenience fields). + metadata_filters: Optional[Dict[str, Any]] = None + if query.metadata_filters or tags or query.status: + metadata_filters = dict(query.metadata_filters or {}) + if tags: + metadata_filters.setdefault("tags", tags) + if query.status: + metadata_filters.setdefault("status", query.status) + + prepared = PreparedSearchQuery( + search_text=search_text, + permalink=query.permalink, + permalink_match=query.permalink_match, + title=query.title, + note_types=( + [normalize_note_type(note_type) for note_type in query.note_types] + if query.note_types + else None + ), + search_item_types=query.entity_types, + categories=query.categories, + after_date=after_date, + metadata_filters=metadata_filters, + file_path_prefix=query.file_path_prefix, + temporal=build_temporal_filter(query), + retrieval_mode=query.retrieval_mode or SearchRetrievalMode.FTS, + min_similarity=query.min_similarity, + ) + + has_criteria = bool( + prepared.search_text + or prepared.permalink + or prepared.permalink_match + or prepared.title + or prepared.note_types + or prepared.search_item_types + or prepared.categories + or prepared.after_date + or prepared.metadata_filters + # Normalized by SearchQuery, so only a real subtree reaches here. + or prepared.file_path_prefix + or prepared.temporal + ) + if not has_criteria: + logger.debug("no criteria passed to query") + return None + return prepared + + +async def include_legacy_note_type_spellings( + session_maker: async_sessionmaker[AsyncSession], + scope: ProjectScope, + prepared: PreparedSearchQuery, + *, + session: AsyncSession | None = None, +) -> PreparedSearchQuery: + """Expand canonical note-type filters to the exact legacy spellings stored in ``scope``. + + Search rows written before canonicalization preserve the owning entity's exact + type spelling. Including those spellings alongside the canonical values keeps an + upgrade searchable without requiring an eager full reindex. Only spellings from + projects in scope are read, so a scope cannot learn what another project stores. + """ + if not prepared.note_types: + return prepared + + canonical_note_types = set(prepared.note_types) + async with db.scoped_session(session_maker, session) as active_session: + stored_types = await active_session.scalars( + select(Entity.note_type).where(Entity.project_id.in_(scope.project_ids)).distinct() + ) + compatible_note_types = canonical_note_types | { + stored_type + for stored_type in stored_types.all() + if stored_type and normalize_note_type(stored_type) in canonical_note_types + } + return replace(prepared, note_types=sorted(compatible_note_types)) + + +def relaxed_fts_fallback_eligible( + query: SearchQuery, + search_text: str | None, + retrieval_mode: SearchRetrievalMode, +) -> bool: + """Whether a zero-result strict full-text query may retry with OR-joined terms.""" + if retrieval_mode != SearchRetrievalMode.FTS: + return False + if not search_text or not search_text.strip(): + return False + if '"' in search_text: + return False + if query.has_boolean_operators(): + return False + # Trigger: query has too few safe relaxed terms, explicit numeric identifiers, + # or only terms that would over-broaden under OR. + # Why: the shared helper preserves the old English guard while allowing + # whitespace-separated CJK terms that ASCII tokenization cannot see. + # Outcome: retry only when there is a backend-safe relaxed OR query. + return relaxed_query_words(search_text) is not None + + class SearchService: """Service for search operations. @@ -196,76 +324,7 @@ async def reindex_all(self, background_tasks: Optional[BackgroundTasks] = None) def prepare_query(self, query: SearchQuery) -> PreparedSearchQuery | None: """Normalize a SearchQuery into repository arguments.""" - search_text = query.text - tags = query.tags - - # Support tag: shorthand by mapping to tags filter. - if search_text is not None: - search_text = search_text.strip() or None - if search_text and search_text.lower().startswith("tag:"): - tag_values = re.split(r"[,\s]+", search_text[4:].strip()) - parsed_tags = [t for t in tag_values if t] - if parsed_tags: - tags = parsed_tags - search_text = None - - after_date = ( - ( - query.after_date - if isinstance(query.after_date, datetime) - else parse(query.after_date) - ) - if query.after_date - else None - ) - - # Merge structured metadata filters (explicit + convenience fields). - metadata_filters: Optional[Dict[str, Any]] = None - if query.metadata_filters or tags or query.status: - metadata_filters = dict(query.metadata_filters or {}) - if tags: - metadata_filters.setdefault("tags", tags) - if query.status: - metadata_filters.setdefault("status", query.status) - - prepared = PreparedSearchQuery( - search_text=search_text, - permalink=query.permalink, - permalink_match=query.permalink_match, - title=query.title, - note_types=( - [normalize_note_type(note_type) for note_type in query.note_types] - if query.note_types - else None - ), - search_item_types=query.entity_types, - categories=query.categories, - after_date=after_date, - metadata_filters=metadata_filters, - file_path_prefix=query.file_path_prefix, - temporal=build_temporal_filter(query), - retrieval_mode=query.retrieval_mode or SearchRetrievalMode.FTS, - min_similarity=query.min_similarity, - ) - - has_criteria = bool( - prepared.search_text - or prepared.permalink - or prepared.permalink_match - or prepared.title - or prepared.note_types - or prepared.search_item_types - or prepared.categories - or prepared.after_date - or prepared.metadata_filters - # Normalized by SearchQuery, so only a real subtree reaches here. - or prepared.file_path_prefix - or prepared.temporal - ) - if not has_criteria: - logger.debug("no criteria passed to query") - return None - return prepared + return prepare_search_query(query) @staticmethod def _prepared_has_filters(prepared: PreparedSearchQuery) -> bool: @@ -286,27 +345,12 @@ async def _include_legacy_note_type_spellings( session: AsyncSession | None = None, ) -> PreparedSearchQuery: """Expand canonical note-type filters to exact legacy entity spellings.""" - if not prepared.note_types: - return prepared - - canonical_note_types = set(prepared.note_types) - async with db.scoped_session(self.session_maker, session) as active_session: - stored_types_query = self.entity_repository.select(Entity.note_type).distinct() - stored_types_result = await self.entity_repository.execute_query( - active_session, - stored_types_query, - use_query_options=False, - ) - - # Search rows written before canonicalization preserve the owning entity's - # exact type spelling. Include those spellings alongside canonical values - # so an upgrade remains searchable without requiring an eager full reindex. - compatible_note_types = canonical_note_types | { - stored_type - for stored_type in stored_types_result.scalars().all() - if stored_type and normalize_note_type(stored_type) in canonical_note_types - } - return replace(prepared, note_types=sorted(compatible_note_types)) + return await include_legacy_note_type_spellings( + self.session_maker, + ProjectScope.single(self.repository.project_id), + prepared, + session=session, + ) async def _search_repository( self, @@ -482,20 +526,7 @@ def _is_relaxed_fts_fallback_eligible( retrieval_mode: SearchRetrievalMode, ) -> bool: """Check whether we should run relaxed OR fallback after strict FTS returns empty.""" - if retrieval_mode != SearchRetrievalMode.FTS: - return False - if not search_text or not search_text.strip(): - return False - if '"' in search_text: - return False - if query.has_boolean_operators(): - return False - # Trigger: query has too few safe relaxed terms, explicit numeric identifiers, - # or only terms that would over-broaden under OR. - # Why: the shared helper preserves the old English guard while allowing - # whitespace-separated CJK terms that ASCII tokenization cannot see. - # Outcome: retry only when there is a backend-safe relaxed OR query. - return relaxed_query_words(search_text) is not None + return relaxed_fts_fallback_eligible(query, search_text, retrieval_mode) @staticmethod def _generate_variants(text: str) -> Set[str]: diff --git a/tests/api/v2/test_scoped_search_router.py b/tests/api/v2/test_scoped_search_router.py new file mode 100644 index 000000000..41a8d3668 --- /dev/null +++ b/tests/api/v2/test_scoped_search_router.py @@ -0,0 +1,657 @@ +"""One query over an explicit set of projects, through the reader and the route. + +Only the embedding provider is a double; projects, entities, search rows, vectors, +retrieval, and hydration are real on whichever backend the session is configured +for. Search row ids are database-wide primary keys, so the corpus gives every row a +distinct id the way real data does. +""" + +import re +from dataclasses import dataclass +from datetime import datetime, timezone +from math import sqrt +from typing import Any +from unittest.mock import AsyncMock + +import pytest +from httpx import AsyncClient +from sqlalchemy import event, text +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker + +import basic_memory.repository.search_repository as search_repository_module +from basic_memory import db +from basic_memory.api.v2.utils import to_search_results +from basic_memory.config import BasicMemoryConfig, DatabaseBackend +from basic_memory.models import Entity, MemoryTimeIndex, Project +from basic_memory.repository.postgres_search_repository import PostgresSearchRepository +from basic_memory.repository.search_index_row import SearchIndexRow +from basic_memory.repository.search_repository import create_search_reader +from basic_memory.repository.search_scope import ProjectScope +from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository +from basic_memory.schemas.search import SearchQuery, SearchRetrievalMode +from basic_memory.services.scoped_search_service import ScopedSearchService +from basic_memory.services.search_service import ( + include_legacy_note_type_spellings, + prepare_search_query, +) + +MARKERS = ("nebula", "comet", "quasar") +# Distinct per project, as real observation and relation primary keys are. +OBSERVATION_IDS = (5001, 5002, 5003) +RELATION_IDS = (6001, 6002, 6003) +SEMANTIC_MODES = [SearchRetrievalMode.VECTOR, SearchRetrievalMode.HYBRID] + + +class MarkerEmbeddingProvider: + """Unit vectors from marker words: deterministic, and different texts rank differently.""" + + model_name = "marker" + dimensions = 4 + + def __init__(self) -> None: + self.query_calls = 0 + + @staticmethod + def _vectorize(text: str) -> list[float]: + words = set(re.findall(r"[a-z]+", text.lower())) + axes = [1.0 if marker in words else 0.0 for marker in MARKERS] + axes.append(0.0 if any(axes) else 1.0) + norm = sqrt(sum(axis * axis for axis in axes)) + return [axis / norm for axis in axes] + + async def embed_query(self, text: str) -> list[float]: + self.query_calls += 1 + return self._vectorize(text) + + async def embed_documents(self, texts: list[str]) -> list[list[float]]: + return [self._vectorize(text) for text in texts] + + def runtime_log_attrs(self) -> dict[str, Any]: + return {} + + +@dataclass +class Corpus: + session_maker: async_sessionmaker[AsyncSession] + engine: AsyncEngine + config: BasicMemoryConfig + projects: list[Project] + entities: list[Entity] + provider: MarkerEmbeddingProvider + + def ids(self, count: int) -> list[int]: + return [project.id for project in self.projects[:count]] + + def service(self, project_ids: list[int]) -> ScopedSearchService: + scope = ProjectScope.of(project_ids) + reader = create_search_reader(self.session_maker, scope, self.config) + return ScopedSearchService(self.session_maker, scope, reader) + + +@pytest.fixture +async def corpus( + engine_factory, app_config: BasicMemoryConfig, monkeypatch: pytest.MonkeyPatch +) -> Corpus: + """Three projects with one note each, all about a nebula, embedded for real. + + Project 0 matches the query vector exactly; projects 1 and 2 carry a second + marker so their vectors sit at cosine 0.707. Project 2 spells its note type the + legacy way, so note-type expansion has something to find. + """ + engine, session_maker = engine_factory + app_config.semantic_search_enabled = True + app_config.reranker_enabled = False + app_config.semantic_min_similarity = 0.0 + provider = MarkerEmbeddingProvider() + monkeypatch.setattr( + search_repository_module, "create_embedding_provider", lambda _config: provider + ) + repository_type = ( + PostgresSearchRepository + if app_config.database_backend == DatabaseBackend.POSTGRES + else SQLiteSearchRepository + ) + now = datetime(2026, 9, 1, tzinfo=timezone.utc) + projects: list[Project] = [] + entities: list[Entity] = [] + for index, flavor in enumerate(("", " comet", " quasar")): + note_type = "Note" if index == 2 else "note" + async with db.scoped_session(session_maker) as session: + project = Project( + name=f"scope-{index}", permalink=f"scope-{index}", path=f"/scope-{index}" + ) + session.add(project) + await session.flush() + entity = Entity( + project_id=project.id, + title="shared nebula", + note_type=note_type, + content_type="text/markdown", + permalink="notes/shared", + file_path="notes/shared.md", + entity_metadata={"status": "open" if index == 0 else "closed", "tags": ["scope"]}, + created_at=now, + updated_at=now, + ) + session.add(entity) + await session.commit() + projects.append(project) + entities.append(entity) + + writer = repository_type( + session_maker, project.id, app_config=app_config, embedding_provider=provider + ) + await writer.init_search_index() + content = f"shared nebula{flavor}" + rows = [ + SearchIndexRow( + project_id=project.id, + id=entity.id, + entity_id=entity.id, + type="entity", + file_path=entity.file_path, + permalink=entity.permalink, + title=entity.title, + content_stems=content, + content_snippet=content, + metadata={"note_type": note_type}, + created_at=now, + updated_at=now, + ) + ] + for row_type, row_id in ( + ("observation", OBSERVATION_IDS[index]), + ("relation", RELATION_IDS[index]), + ): + rows.append( + SearchIndexRow( + project_id=project.id, + id=row_id, + entity_id=entity.id, + type=row_type, + file_path=entity.file_path, + permalink=f"notes/shared/{row_type}", + title="shared nebula", + content_stems=content, + content_snippet=f"{content} project {index} {row_type}", + category="fact" if row_type == "observation" else None, + from_id=entity.id if row_type == "relation" else None, + relation_type="relates_to" if row_type == "relation" else None, + created_at=now, + updated_at=now, + ) + ) + await writer.bulk_index_items(rows) + await writer.sync_entity_vectors(entity.id) + return Corpus(session_maker, engine, app_config, projects, entities, provider) + + +async def _mark_pending(corpus: Corpus, project_ids: list[int] | None = None) -> None: + """Take vectors out of play so a hybrid answer proves the lexical channel.""" + async with db.scoped_session(corpus.session_maker) as session: + if project_ids is None: + await session.execute( + text("UPDATE search_vector_chunks SET embedding_status = 'pending'") + ) + else: + for project_id in project_ids: + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'pending' " + "WHERE project_id = :project_id" + ), + {"project_id": project_id}, + ) + await session.commit() + + +async def _add_temporal_assertions(corpus: Corpus, project_ids: list[int]) -> None: + """Project 0's observation is current; every other project's ended in 2025.""" + async with db.scoped_session(corpus.session_maker) as session: + for index, (project, entity) in enumerate(zip(corpus.projects, corpus.entities)): + if project.id not in project_ids: + continue + current = index == 0 + session.add( + MemoryTimeIndex( + project_id=project.id, + entity_id=entity.id, + source_type="observation", + source_id=OBSERVATION_IDS[index], + time_kind="effective", + range_axis="date", + lower_value="2026-09-01" if current else "2025-01-01", + upper_value=None if current else "2025-02-01", + lower_inclusive=True, + upper_inclusive=False, + is_empty=False, + extractor="test", + source_text=( + "@effective[2026-09-01,)" + if current + else "@effective[2025-01-01,2025-02-01)" + ), + ) + ) + await session.commit() + + +# --- The reader over a scope --- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", list(SearchRetrievalMode)) +@pytest.mark.parametrize("scope_size", [0, 1, 2]) +async def test_one_reader_answers_the_whole_scope( + corpus: Corpus, mode: SearchRetrievalMode, scope_size: int +) -> None: + """Every project in scope answers from one pipeline, and the pipeline only reads.""" + service = corpus.service(corpus.ids(scope_size)) + query = SearchQuery(text="nebula", retrieval_mode=mode) + statements: list[str] = [] + + def record(_conn, _cursor, statement, _params, _context, _many): + statements.append(statement) + + event.listen(corpus.engine.sync_engine, "before_cursor_execute", record) + try: + rows = await service.search(query, limit=100, offset=0) + finally: + event.remove(corpus.engine.sync_engine, "before_cursor_execute", record) + + assert len(rows) == scope_size * 3 + assert {row.project_id for row in rows} == set(corpus.ids(scope_size)) + assert len({(row.type, row.id) for row in rows}) == len(rows) + # One embedding per search, and none at all for an empty scope. + assert corpus.provider.query_calls == int(scope_size > 0 and mode != SearchRetrievalMode.FTS) + mutations = [ + sql + for sql in statements + if sql.lstrip().upper().startswith(("INSERT", "UPDATE", "DELETE", "DROP")) + ] + assert mutations == [] + if mode == SearchRetrievalMode.FTS: + assert await service.count(query) == scope_size * 3 + else: + assert all(row.matched_chunk_text for row in rows) + with pytest.raises(ValueError, match="Exact counts"): + await service.count(query) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", list(SearchRetrievalMode)) +async def test_pages_are_a_stable_slice_of_the_complete_order( + corpus: Corpus, mode: SearchRetrievalMode +) -> None: + service = corpus.service(corpus.ids(2)) + query = SearchQuery(text="nebula", retrieval_mode=mode) + + complete = await service.search(query, limit=100, offset=0) + pages = [ + row + for offset in range(0, len(complete), 2) + for row in await service.search(query, limit=2, offset=offset) + ] + + assert [(r.project_id, r.type, r.id, r.score) for r in pages] == [ + (r.project_id, r.type, r.id, r.score) for r in complete + ] + assert await service.search(query, limit=10, offset=100) == [] + + +@pytest.mark.asyncio +async def test_vectors_outside_the_current_manifest_are_not_served(corpus: Corpus) -> None: + """Pending, re-modelled, and stale-source chunks are invisible across the whole scope.""" + ids = corpus.ids(3) + async with db.scoped_session(corpus.session_maker) as session: + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'pending' WHERE project_id = :p" + ), + {"p": ids[0]}, + ) + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_model = 'other-model' WHERE project_id = :p" + ), + {"p": ids[1]}, + ) + await session.execute( + text("UPDATE search_vector_chunks SET source_hash = 'stale' WHERE project_id = :p"), + {"p": ids[2]}, + ) + await session.commit() + + query = SearchQuery(text="nebula", retrieval_mode=SearchRetrievalMode.VECTOR) + assert await corpus.service(ids).search(query, limit=100, offset=0) == [] + + +@pytest.mark.asyncio +async def test_hybrid_keeps_lexical_only_rows_below_the_threshold(corpus: Corpus) -> None: + ids = corpus.ids(2) + await _mark_pending(corpus, [ids[1]]) + query = SearchQuery(text="nebula", retrieval_mode=SearchRetrievalMode.HYBRID, min_similarity=1) + + rows = await corpus.service(ids).search(query, limit=100, offset=0) + + assert len(rows) == 6 + assert all(row.matched_chunk_text is None for row in rows if row.project_id == ids[1]) + assert max(row.score or 0.0 for row in rows) == pytest.approx(1.3) + assert all(1.0 <= (row.score or 0.0) <= 1.3 + 1e-6 for row in rows if row.project_id == ids[0]) + assert all(0.0 <= (row.score or 0.0) <= 1.0 for row in rows if row.project_id == ids[1]) + + +@pytest.mark.asyncio +async def test_legacy_note_type_spellings_expand_only_within_scope(corpus: Corpus) -> None: + """A scope learns the spellings its own projects store, and nothing about others.""" + prepared = prepare_search_query(SearchQuery(text="nebula", note_types=["note"])) + assert prepared is not None + canonical_only = ProjectScope.of(corpus.ids(2)) + legacy_project = ProjectScope.single(corpus.projects[2].id) + + within_canonical = await include_legacy_note_type_spellings( + corpus.session_maker, canonical_only, prepared + ) + within_legacy = await include_legacy_note_type_spellings( + corpus.session_maker, legacy_project, prepared + ) + + assert within_canonical.note_types == ["note"] + assert within_legacy.note_types == ["Note", "note"] + # And the expanded filter finds the legacy rows when that project is in scope. + query = SearchQuery(text="nebula", note_types=["note"]) + assert len(await corpus.service(corpus.ids(3)).search(query, limit=100, offset=0)) == 9 + assert ( + len(await corpus.service([corpus.projects[2].id]).search(query, limit=100, offset=0)) == 3 + ) + + +@pytest.mark.asyncio +async def test_hydration_reads_only_projects_in_scope(corpus: Corpus) -> None: + ids = corpus.ids(2) + await _add_temporal_assertions(corpus, ids) + service = corpus.service([ids[0]]) + all_entity_ids = [entity.id for entity in corpus.entities] + + entities = await service.get_entities_by_id(all_entity_ids) + # Search before opening the hydration session: the test pool holds one connection. + rows = await service.search(SearchQuery(text="nebula"), limit=100, offset=0) + async with db.scoped_session(corpus.session_maker) as session: + assertions = await service.find_for_sources( + session, [("observation", source_id) for source_id in OBSERVATION_IDS] + ) + external_ids = await service.project_external_ids(session, rows) + + assert [entity.id for entity in entities] == [corpus.entities[0].id] + assert [row.project_id for row in assertions] == [ids[0]] + assert external_ids == {ids[0]: corpus.projects[0].external_id} + # An empty scope hydrates nothing, and an empty page names no projects. + unscoped = corpus.service([]) + assert await unscoped.get_entities_by_id(all_entity_ids) == [] + async with db.scoped_session(corpus.session_maker) as session: + assert await unscoped.find_for_sources(session, [("observation", OBSERVATION_IDS[0])]) == [] + assert await service.project_external_ids(session, []) == {} + + +# --- The route --- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", list(SearchRetrievalMode)) +async def test_filters_narrow_the_scope_and_results_carry_project_identity( + corpus: Corpus, client: AsyncClient, mode: SearchRetrievalMode +) -> None: + ids = corpus.ids(2) + + response = await client.request( + "QUERY", + "/v2/search/", + json={ + "project_ids": ids, + "text": "nebula", + "retrieval_mode": mode.value, + "metadata_filters": {"status": "open"}, + "note_types": ["note"], + "file_path_prefix": "notes", + "entity_types": ["observation"], + "categories": ["fact"], + }, + ) + + assert response.status_code == 200, response.text + data = response.json() + assert len(data["results"]) == 1 + row = data["results"][0] + assert row["project_id"] == ids[0] + assert row["project_external_id"] == corpus.projects[0].external_id + assert row["external_id"] == corpus.entities[0].external_id + assert row["observation_id"] == OBSERVATION_IDS[0] + assert row["entity"] == "notes/shared" + assert data["total_is_exact"] == (mode == SearchRetrievalMode.FTS) + assert data["total"] == int(mode == SearchRetrievalMode.FTS) + assert not data["has_more"] + assert response.headers["cache-control"] == "no-store" + assert response.headers["accept-query"] == "application/json" + + +@pytest.mark.asyncio +async def test_scope_shapes_the_answer(corpus: Corpus, client: AsyncClient) -> None: + for ids, count in [(corpus.ids(3), 9), ([corpus.projects[1].id], 3), ([], 0)]: + response = await client.request( + "QUERY", "/v2/search/", json={"project_ids": ids, "text": "nebula"} + ) + assert response.status_code == 200, response.text + assert response.json()["total"] == count + assert {r["project_id"] for r in response.json()["results"]} <= set(ids) + + +@pytest.mark.asyncio +async def test_post_is_the_documented_twin_of_query( + corpus: Corpus, client: AsyncClient, app +) -> None: + body = {"project_ids": corpus.ids(2), "text": "nebula"} + + posted = await client.post("/v2/search/", json=body) + queried = await client.request("QUERY", "/v2/search/", json=body) + + assert posted.status_code == 200, posted.text + assert posted.json() == queried.json() + assert set(app.openapi()["paths"]["/v2/search/"]) == {"post"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope", [None, [0], [True], ["1"]]) +async def test_route_rejects_an_unspecified_or_invalid_scope( + client: AsyncClient, scope: object +) -> None: + response = await client.request( + "QUERY", "/v2/search/", json={"text": "nebula", "project_ids": scope} + ) + assert response.status_code == 422 + response = await client.request("QUERY", "/v2/search/", json={"text": "nebula"}) + assert response.status_code == 422 + response = await client.request( + "QUERY", "/v2/search/", params={"page": 0}, json={"text": "nebula", "project_ids": [1]} + ) + assert response.status_code == 422 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", list(SearchRetrievalMode)) +async def test_temporal_filter_narrows_and_explains_within_scope( + corpus: Corpus, client: AsyncClient, mode: SearchRetrievalMode +) -> None: + ids = corpus.ids(2) + await _add_temporal_assertions(corpus, ids) + + response = await client.request( + "QUERY", + "/v2/search/", + json={ + "project_ids": ids, + "text": "nebula", + "retrieval_mode": mode.value, + "valid_at": "2026-09-14", + "note_types": ["note"], + }, + ) + + assert response.status_code == 200, response.text + payload = response.json() + assert payload["temporal_applied"] is True + assert len(payload["results"]) == 1 + assert payload["results"][0]["project_id"] == ids[0] + assert payload["results"][0]["temporal"][0]["source_text"] == "@effective[2026-09-01,)" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", list(SearchRetrievalMode)) +async def test_api_pagination_and_empty_later_page( + corpus: Corpus, client: AsyncClient, mode: SearchRetrievalMode +) -> None: + query = {"project_ids": corpus.ids(2), "text": "nebula", "retrieval_mode": mode.value} + for page, count, has_more in [(1, 2, True), (3, 2, False), (4, 0, False)]: + response = await client.request( + "QUERY", "/v2/search/", params={"page": page, "page_size": 2}, json=query + ) + assert response.status_code == 200, response.text + assert len(response.json()["results"]) == count + assert response.json()["has_more"] == has_more + + +@pytest.mark.asyncio +async def test_semantic_modes_need_the_semantic_stack(corpus: Corpus, client: AsyncClient) -> None: + corpus.config.semantic_search_enabled = False + scope = {"project_ids": [corpus.projects[0].id]} + + for mode in SEMANTIC_MODES: + response = await client.request( + "QUERY", "/v2/search/", json={**scope, "text": "nebula", "retrieval_mode": mode.value} + ) + assert response.status_code == 400, response.text + assert "disabled" in response.json()["detail"] + assert corpus.provider.query_calls == 0 + + no_criteria = await client.request("QUERY", "/v2/search/", json=scope) + assert no_criteria.status_code == 200 and not no_criteria.json()["results"] + lexical = await client.request("QUERY", "/v2/search/", json={**scope, "text": "nebula"}) + assert lexical.status_code == 200, lexical.text + assert lexical.json()["total"] == 3 + + +@pytest.mark.asyncio +async def test_project_route_results_carry_project_identity( + corpus: Corpus, client: AsyncClient +) -> None: + project = corpus.projects[0] + + response = await client.post( + f"/v2/projects/{project.external_id}/search/", json={"text": "nebula"} + ) + + assert response.status_code == 200, response.text + results = response.json()["results"] + assert len(results) == 3 + assert {row["project_id"] for row in results} == {project.id} + assert {row["project_external_id"] for row in results} == {project.external_id} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["fts", "hybrid"]) +@pytest.mark.parametrize("query", ["Did nebula go hiking at sunrise?", "foo None: + ids = corpus.ids(2) + # Remove the vector channel so a hybrid success proves lexical recovery. + await _mark_pending(corpus) + body = {"project_ids": ids, "text": query, "retrieval_mode": mode} + + complete = await client.request("QUERY", "/v2/search/", json=body) + + assert complete.status_code == 200, complete.text + expected = complete.json()["results"] + assert len(expected) == 6 + assert {row["project_id"] for row in expected} == set(ids) + if mode == "fts": + assert complete.json()["total"] == 6 + pages = [] + for page in range(1, 5): + response = await client.request( + "QUERY", "/v2/search/", json=body, params={"page": page, "page_size": 2} + ) + assert response.status_code == 200, response.text + pages.extend(response.json()["results"]) + assert response.json()["has_more"] == (page < 3) + assert pages == expected + assert corpus.provider.query_calls == (5 if mode == "hybrid" else 0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["fts", "hybrid"]) +async def test_relaxation_keeps_strict_matches_on_deep_pages( + corpus: Corpus, client: AsyncClient, mode: str +) -> None: + # Three terms permit relaxation, but strict matches anywhere in the selected + # scope must suppress it even after the final page. + await _mark_pending(corpus) + body = { + "project_ids": corpus.ids(2), + "text": "shared nebula observation", + "retrieval_mode": mode, + "min_similarity": 1, + } + + first = await client.request("QUERY", "/v2/search/", json=body) + later = await client.request("QUERY", "/v2/search/", json=body, params={"page": 2}) + + assert first.status_code == 200, first.text + assert len(first.json()["results"]) == 2 + assert later.status_code == 200, later.text + assert later.json()["results"] == [] + if mode == "fts": + assert later.json()["total"] == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("query", ["foo None: + response = await client.request( + "QUERY", + "/v2/search/", + json={"project_ids": [corpus.projects[0].id], "text": query, "retrieval_mode": mode}, + ) + + assert response.status_code == 200, response.text + assert len(response.json()["results"]) == (3 if mode == "hybrid" else 0) + if mode == "fts": + assert response.json()["total"] == 0 + assert corpus.provider.query_calls == int(mode == "hybrid") + + +# --- Hydration shaping --- + + +@pytest.mark.asyncio +async def test_results_name_their_project_when_the_caller_knows_it() -> None: + now = datetime(2026, 9, 1, tzinfo=timezone.utc) + row = SearchIndexRow( + project_id=7, + id=1, + type="entity", + title="Alpha", + file_path="alpha.md", + created_at=now, + updated_at=now, + ) + lookup = AsyncMock() + lookup.get_entities_by_id = AsyncMock(return_value=[]) + + named = await to_search_results(lookup, [row], project_external_ids={7: "project-7"}) + unnamed = await to_search_results(lookup, [row]) + + assert (named[0].project_id, named[0].project_external_id) == (7, "project-7") + assert (unnamed[0].project_id, unnamed[0].project_external_id) == (7, None) diff --git a/tests/api/v2/test_search_router_telemetry.py b/tests/api/v2/test_search_router_telemetry.py index d1eec98ba..eeb459975 100644 --- a/tests/api/v2/test_search_router_telemetry.py +++ b/tests/api/v2/test_search_router_telemetry.py @@ -32,7 +32,9 @@ def fake_span(name: str, **attrs): operations.append((name, attrs)) yield - async def fake_to_search_results(entity_service, results, *, temporal_by_source=None): + async def fake_to_search_results( + entity_service, results, *, temporal_by_source=None, project_external_ids=None + ): return [] monkeypatch.setattr(logfire, "span", fake_span) @@ -48,6 +50,7 @@ async def fake_to_search_results(entity_service, results, *, temporal_by_source= session_maker=object(), read_cache=None, response=http_response, + internal_project_id=1, project_id="11111111-1111-1111-1111-111111111111", page=2, page_size=5, diff --git a/tests/repository/test_search_reader.py b/tests/repository/test_search_reader.py index cccf2668b..de29903cb 100644 --- a/tests/repository/test_search_reader.py +++ b/tests/repository/test_search_reader.py @@ -281,3 +281,20 @@ async def test_hybrid_trace_records_the_stable_pool_refetch(): assert trace.stable_pool_refetched is True assert [row.id for row in page] == [4] + + +@pytest.mark.asyncio +async def test_empty_scope_answers_nothing_without_embedding_or_lexical_work() -> None: + """An empty scope admits no rows, so neither leg runs and the query is never embedded.""" + fts = FakeFts([FakeRow(id=1)]) + embed_query = AsyncMock(side_effect=AssertionError("must not embed for an empty scope")) + semantic = SemanticSearch( + cast(Any, None), ProjectScope.of([]), fts, fake_vector_retrieval(embed_query=embed_query) + ) + vector_query = PreparedSearchQuery( + search_text="auth", retrieval_mode=SearchRetrievalMode.VECTOR + ) + + assert await semantic.vector_only(vector_query, limit=10, offset=0) == [] + assert await semantic.hybrid(HYBRID_QUERY, limit=10, offset=0) == [] + assert fts.queries == [] diff --git a/tests/repository/test_search_reader_factory.py b/tests/repository/test_search_reader_factory.py new file mode 100644 index 000000000..88c96de26 --- /dev/null +++ b/tests/repository/test_search_reader_factory.py @@ -0,0 +1,115 @@ +"""Composing a SearchReader over a scope without a project repository.""" + +from typing import Any +from unittest.mock import MagicMock + +import pytest + +import basic_memory.repository.search_repository as search_repository_module +from basic_memory.config import BasicMemoryConfig, DatabaseBackend +from basic_memory.repository.postgres_search_query import PostgresFts +from basic_memory.repository.search_reader import Reranking +from basic_memory.repository.search_repository import create_search_reader +from basic_memory.repository.search_scope import ProjectScope +from basic_memory.repository.semantic_vector_index_factory import semantic_embedding_identity +from basic_memory.repository.sqlite_search_query import SQLiteFts + +SCOPE = ProjectScope.of([3, 1]) + + +class _StubEmbeddingProvider: + model_name = "stub" + dimensions = 4 + + async def embed_query(self, text: str) -> list[float]: + return [1.0, 0.0, 0.0, 0.0] + + async def embed_documents(self, texts: list[str]) -> list[list[float]]: + return [[1.0, 0.0, 0.0, 0.0] for _ in texts] + + def runtime_log_attrs(self) -> dict[str, Any]: + return {} + + +class _StubReranker: + model_name = "stub-reranker" + + async def rerank(self, query: str, documents: list[str]) -> list[float]: + return [0.5 for _ in documents] + + def runtime_log_attrs(self) -> dict[str, Any]: + return {} + + +def _config(backend: DatabaseBackend, **overrides: object) -> BasicMemoryConfig: + return BasicMemoryConfig( + env="test", + projects={"test-project": "/tmp/test"}, + default_project="test-project", + database_backend=backend, + **overrides, + ) + + +@pytest.mark.parametrize( + ("backend", "fts_type"), + [(DatabaseBackend.SQLITE, SQLiteFts), (DatabaseBackend.POSTGRES, PostgresFts)], +) +def test_reader_without_semantic_search_is_full_text_only(monkeypatch, backend, fts_type): + """Disabled semantic search never resolves a provider, and the engine picks the backend.""" + monkeypatch.setattr( + search_repository_module, + "create_embedding_provider", + lambda _config: pytest.fail("a full-text reader must not load an embedding provider"), + ) + + reader = create_search_reader( + MagicMock(), SCOPE, _config(backend, semantic_search_enabled=False) + ) + + assert reader.scope == SCOPE + assert reader.semantic is None + assert isinstance(reader.fts, fts_type) + + +@pytest.mark.parametrize("reranker", [None, _StubReranker()]) +def test_reader_with_semantic_search_composes_the_shared_stack(monkeypatch, reranker): + """The reader gets the same provider, adapter, and reranker a project repository would.""" + provider = _StubEmbeddingProvider() + index = MagicMock() + captured: dict[str, Any] = {} + + def fake_create_index(**kwargs: Any) -> tuple[str, Any]: + captured.update(kwargs) + return "sqlite-vec", index + + monkeypatch.setattr(search_repository_module, "create_embedding_provider", lambda _c: provider) + monkeypatch.setattr(search_repository_module, "create_semantic_vector_index", fake_create_index) + monkeypatch.setattr(search_repository_module, "create_rerank_provider", lambda _c: reranker) + config = _config( + DatabaseBackend.SQLITE, + semantic_search_enabled=True, + semantic_vector_k=7, + semantic_min_similarity=0.25, + reranker_candidates=9, + reranker_max_document_chars=123, + ) + + reader = create_search_reader(MagicMock(), SCOPE, config) + + semantic = reader.semantic + assert semantic is not None + assert semantic.scope == SCOPE and semantic.fts is reader.fts + assert semantic.vector.index is index + assert semantic.vector.index_name == "sqlite-vec" + assert semantic.vector.embedding_provider is provider + assert semantic.vector.embedding_model == semantic_embedding_identity(provider) + assert (semantic.vector.vector_k, semantic.vector.min_similarity) == (7, 0.25) + # The adapter is built for the database, not for a project. + assert captured["database_backend"] == DatabaseBackend.SQLITE + assert captured["embedding_provider"] is provider + assert "project_id" not in captured + if reranker is None: + assert semantic.rerank is None + else: + assert semantic.rerank == Reranking(provider=reranker, candidates=9, max_document_chars=123) diff --git a/tests/repository/test_sqlite_vec_index.py b/tests/repository/test_sqlite_vec_index.py new file mode 100644 index 000000000..b5c613a5d --- /dev/null +++ b/tests/repository/test_sqlite_vec_index.py @@ -0,0 +1,41 @@ +"""sqlite-vec adapter behavior that does not need a database.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +from sqlalchemy.exc import OperationalError as SAOperationalError + +from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError +from basic_memory.repository.semantic_vector_index import VectorIndexScope +from basic_memory.repository.sqlite_vec_index import SQLiteVecIndex + + +@pytest.mark.asyncio +async def test_missing_extension_support_is_a_dependency_error() -> None: + """A Python whose sqlite3 cannot load extensions gets the keyword-only fallback error. + + aiosqlite exposes ``enable_load_extension`` even when the wrapped connection does + not, so the attribute probe passes and only the call reveals it (#711). The adapter + must raise the same typed error the repository does, or a scoped vector search on + such a host answers 500 instead of an actionable 400. + """ + index = SQLiteVecIndex( + MagicMock(), + VectorIndexScope(namespace="test", embedding_identity="test", dimensions=4), + ) + driver = SimpleNamespace( + enable_load_extension=AsyncMock(side_effect=AttributeError("enable_load_extension")) + ) + connection = AsyncMock() + connection.get_raw_connection = AsyncMock( + return_value=SimpleNamespace(driver_connection=driver) + ) + session = AsyncMock() + session.execute = AsyncMock( + side_effect=SAOperationalError("SELECT vec_version()", {}, Exception("no such function")) + ) + session.connection = AsyncMock(return_value=connection) + + with pytest.raises(SemanticDependenciesMissingError, match="extension loading"): + await index._ensure_loaded(session) From 5ab7f1febee3311c6782777cd2851b002cd3be45 Mon Sep 17 00:00:00 2001 From: phernandez Date: Tue, 15 Sep 2026 17:35:59 -0500 Subject: [PATCH 6/8] feat(mcp): search all projects with one scoped query per database `search_notes(search_all_projects=True)` used to run one project search per project and merge the pages client-side. It now groups the accessible projects by the database they live in (the local database, or one cloud workspace each) and runs one `QUERY /v2/search/` per group through the new `ScopedSearchClient`, passing the group's internal project ids. The server ranks the union in one query and attributes each hit with `project_external_id`; the tool qualifies permalinks and paths from that attribution rather than from which request a result came back on, so a name collision across workspaces can no longer route a hit to the wrong project. A new `projects` parameter searches a chosen subset by name or external id. It shares the grouping and merge with the all-projects path; an unknown name is an error before any query runs. Cross-database results are still merged by score and sliced to the requested page, a failing database is skipped with a warning and an inexact total, a retryable outage fails the merged page, and a server too old to attribute its results raises a version-skew error like the existing valid-time skew check. `list_memory_projects` now carries the internal `id` in its merged project rows, which is what the scoped route addresses projects by. The temporal confirmation check moves to a module function shared by both search clients. Refs #1558 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_019YW9ysxugGGBCNEGzsxtFV Signed-off-by: phernandez --- CHANGELOG.md | 9 + src/basic_memory/mcp/clients/__init__.py | 3 +- src/basic_memory/mcp/clients/search.py | 89 ++- .../mcp/tools/project_management.py | 3 + src/basic_memory/mcp/tools/search.py | 623 ++++++++++------- tests/mcp/test_tool_contracts.py | 1 + tests/mcp/test_tool_search_temporal.py | 36 +- tests/mcp/tools/test_cjk_search_guidance.py | 104 ++- .../tools/test_search_notes_multi_project.py | 661 ++++++++++-------- ...est_search_notes_multi_project_temporal.py | 90 ++- 10 files changed, 949 insertions(+), 670 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a53634749..cc000c6df 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,15 @@ hydrated only from projects in scope. Every search result, on both routes, now carries `project_id` and `project_external_id`. +- **#1558**: `search_notes(search_all_projects=True)` runs one scoped query per database + instead of one search per project, so a local vault with many projects is one + query, and each cloud workspace is one query, with results attributed to their + project by the server rather than by which request they came back on. A new + `projects` parameter searches a chosen subset by name or external id; an unknown + name is an error. Cross-database results are still merged by score, one failing + database is skipped with a warning and an inexact total, and a retryable outage + or a server too old to attribute its results fails the whole page. + - **#1512**: Word, PowerPoint, and CSV files get the same sidecar Markdown note a PDF gets. `bm import document ` indexes the project, extracts the file, and writes `..md` next to it plus a run note under diff --git a/src/basic_memory/mcp/clients/__init__.py b/src/basic_memory/mcp/clients/__init__.py index 303ff9f67..28190add2 100644 --- a/src/basic_memory/mcp/clients/__init__.py +++ b/src/basic_memory/mcp/clients/__init__.py @@ -12,7 +12,7 @@ """ from basic_memory.mcp.clients.knowledge import KnowledgeClient -from basic_memory.mcp.clients.search import SearchClient +from basic_memory.mcp.clients.search import ScopedSearchClient, SearchClient from basic_memory.mcp.clients.memory import MemoryClient from basic_memory.mcp.clients.directory import DirectoryClient from basic_memory.mcp.clients.resource import ResourceClient @@ -22,6 +22,7 @@ __all__ = [ "KnowledgeClient", + "ScopedSearchClient", "SearchClient", "MemoryClient", "DirectoryClient", diff --git a/src/basic_memory/mcp/clients/search.py b/src/basic_memory/mcp/clients/search.py index 08886178f..932e964e5 100644 --- a/src/basic_memory/mcp/clients/search.py +++ b/src/basic_memory/mcp/clients/search.py @@ -1,8 +1,10 @@ -"""Typed client for search API operations. +"""Typed clients for search API operations. -Encapsulates all /v2/projects/{project_id}/search/* endpoints. +``SearchClient`` covers /v2/projects/{project_id}/search/*; ``ScopedSearchClient`` +covers QUERY /v2/search/, one search over an explicit set of projects. """ +from collections.abc import Mapping, Sequence from typing import Any from httpx import AsyncClient @@ -19,6 +21,71 @@ _TEMPORAL_QUERY_FIELDS = ("valid_at", "valid_overlaps", "time_kind") +def _confirm_temporal_filter_applied(query: Mapping[str, Any], payload: Mapping[str, Any]) -> None: + """Refuse a response that does not confirm a requested valid-time filter ran. + + Trigger: this request carried a valid-time filter but the response does not + confirm the server ran it. + Why: SearchQuery ignores unknown fields, so a server predating SPEC-82 accepts + the request and returns results that look filtered. A valid-time query + excludes undated sources; unfiltered results include them, and the caller + would have no way to tell. + Outcome: fail loudly instead of returning a wrong answer that reads as right. + """ + if any(query.get(field) for field in _TEMPORAL_QUERY_FIELDS) and ( + payload.get("temporal_applied") is not True + ): + raise ValueError( + "The search API did not apply the requested valid-time filter " + "(no temporal_applied confirmation in the response). The server is " + "likely older than this client; upgrade it or drop valid_at / " + "valid_overlaps / time_kind from the query." + ) + + +class ScopedSearchClient: + """Typed client for ``QUERY /v2/search/``: one search over an explicit set of projects. + + ``project_ids`` are the internal ids the project list reports; every id must + belong to the database the HTTP client is routed to. An empty set answers no rows. + """ + + def __init__(self, http_client: AsyncClient): + self.http_client = http_client + + async def search( + self, + query: dict[str, Any], + *, + project_ids: Sequence[int], + page: int = 1, + page_size: int = 10, + ) -> SearchResponse: + """Search the named projects; results carry ``project_id`` and ``project_external_id``.""" + from basic_memory.mcp.tools.utils import call_query + + with logfire.span( + "mcp.client.search.search_scope", + client_name="search", + operation="search_scope", + project_count=len(project_ids), + page=page, + page_size=page_size, + ): + response = await call_query( + self.http_client, + "/v2/search/", + json={**query, "project_ids": list(project_ids)}, + params={"page": page, "page_size": page_size}, + client_name="search", + operation="search_scope", + path_template="/v2/search/", + ) + payload = response.json() + _confirm_temporal_filter_applied(query, payload) + return SearchResponse.model_validate(payload) + + class SearchClient: """Typed client for search operations. @@ -92,21 +159,5 @@ async def search( retrieval_mode = query.get("retrieval_mode", SearchRetrievalMode.FTS) payload["total_is_exact"] = retrieval_mode == SearchRetrievalMode.FTS - # Trigger: this request carried a valid-time filter but the response does not - # confirm the server ran it. - # Why: SearchQuery ignores unknown fields, so a server predating SPEC-82 accepts - # the request and returns results that look filtered. A valid-time query - # excludes undated sources; unfiltered results include them, and the caller - # would have no way to tell. - # Outcome: fail loudly instead of returning a wrong answer that reads as right. - if any(query.get(field) for field in _TEMPORAL_QUERY_FIELDS) and ( - payload.get("temporal_applied") is not True - ): - raise ValueError( - "The search API did not apply the requested valid-time filter " - "(no temporal_applied confirmation in the response). The server is " - "likely older than this client; upgrade it or drop valid_at / " - "valid_overlaps / time_kind from the query." - ) - + _confirm_temporal_filter_applied(query, payload) return SearchResponse.model_validate(payload) diff --git a/src/basic_memory/mcp/tools/project_management.py b/src/basic_memory/mcp/tools/project_management.py index e59bbee78..b9bd53407 100644 --- a/src/basic_memory/mcp/tools/project_management.py +++ b/src/basic_memory/mcp/tools/project_management.py @@ -124,6 +124,7 @@ def _merge_projects( { "name": name, "external_id": external_id, + "id": proj.id if proj else None, "path": path, "local_path": local_path, "cloud_path": cloud_path, @@ -258,6 +259,7 @@ def _merge_workspace_projects( { "name": cloud_proj.name, "external_id": cloud_proj.external_id, + "id": cloud_proj.id, "path": local_path or cloud_path, "local_path": local_path, "cloud_path": cloud_path, @@ -283,6 +285,7 @@ def _merge_workspace_projects( { "name": project.name, "external_id": project.external_id, + "id": project.id, "path": project.path, "local_path": project.path, "cloud_path": None, diff --git a/src/basic_memory/mcp/tools/search.py b/src/basic_memory/mcp/tools/search.py index 75f75f409..e02da891f 100644 --- a/src/basic_memory/mcp/tools/search.py +++ b/src/basic_memory/mcp/tools/search.py @@ -1,6 +1,7 @@ """Search tools for Basic Memory MCP server.""" import re +from dataclasses import dataclass from textwrap import dedent from typing import Annotated, List, Optional, Dict, Any, Literal, cast from uuid import UUID @@ -11,7 +12,7 @@ from fastmcp import Context from pydantic import AliasChoices, BeforeValidator, Field -from basic_memory.config import ConfigManager, has_cloud_credentials +from basic_memory.config import ConfigManager from basic_memory.utils import ( build_canonical_permalink, coerce_dict, @@ -19,11 +20,7 @@ parse_tags, strict_search_tags, ) -from basic_memory.mcp.async_client import ( - _explicit_routing, - _force_local_mode, - is_factory_mode, -) +from basic_memory.mcp.async_client import get_client from basic_memory.mcp.container import get_container from basic_memory.mcp.index_readiness import project_index_required from basic_memory.mcp.project_context import ( @@ -33,6 +30,7 @@ ) from basic_memory.mcp.server import mcp from basic_memory.schemas.base import normalize_note_type +from basic_memory.schemas.project_info import ProjectItem from basic_memory.schemas.search import ( SearchItemType, SearchQuery, @@ -43,6 +41,118 @@ from basic_memory.temporal import TemporalQualifierError, parse_temporal_filter _SERVICE_UNAVAILABLE_HEADING = "# Search Failed - Service Temporarily Unavailable" +_NO_SEARCH_CRITERIA_MESSAGE = ( + "# No Search Criteria\n\n" + "Please provide at least one of: `query`, `metadata_filters`, " + "`tags`, `status`, `note_types`, `entity_types`, `categories`, " + "`after_date`, `valid_at`, `valid_overlaps`, or `time_kind`." +) +# Alias common column/model names to their frontmatter key equivalents. Users often +# pass "note_type" (the entity model column) when the frontmatter field is "type". +_METADATA_KEY_ALIASES = {"note_type": "type"} +_VALID_SEARCH_TYPES = ("hybrid", "permalink", "semantic", "text", "title", "vector") + + +def _build_search_query( + *, + query: str | None, + search_type: str, + note_types: list[str], + entity_types: list[str], + categories: list[str], + after_date: str | None, + metadata_filters: dict[str, Any] | None, + tags: list[str] | None, + status: str | None, + min_similarity: float | None, + valid_at: str | None, + valid_overlaps: str | None, + time_kind: str | None, +) -> SearchQuery | None: + """Map tool parameters onto one ``SearchQuery``; ``None`` when nothing narrows the search. + + Shared by the project search and the all-projects search so both ask the API + the same question for the same parameters. + """ + search_query = SearchQuery() + + # Only map search_type to query fields when there is an actual query string. + # When query is None/empty, skip the search mode block: filters-only path. + effective_query = (query or "").strip() + if effective_query: + if search_type == "text": + search_query.text = effective_query + search_query.retrieval_mode = SearchRetrievalMode.FTS + elif search_type in ("vector", "semantic"): + search_query.text = effective_query + search_query.retrieval_mode = SearchRetrievalMode.VECTOR + elif search_type == "hybrid": + search_query.text = effective_query + search_query.retrieval_mode = SearchRetrievalMode.HYBRID + elif search_type == "title": + search_query.title = effective_query + elif search_type == "permalink" and "*" in effective_query: + search_query.permalink_match = effective_query + elif search_type == "permalink": + search_query.permalink = effective_query + else: + raise ValueError( + f"Invalid search_type '{search_type}'. " + f"Valid options: {', '.join(_VALID_SEARCH_TYPES)}" + ) + + # Add optional filters if provided (empty lists are treated as no filter) + if entity_types: + search_query.entity_types = [SearchItemType(t) for t in entity_types] + if categories: + search_query.categories = categories + if note_types: + search_query.note_types = note_types + if after_date: + search_query.after_date = after_date + if metadata_filters: + search_query.metadata_filters = { + _METADATA_KEY_ALIASES.get(key, key): value for key, value in metadata_filters.items() + } + if tags: + search_query.tags = tags + if status: + search_query.status = status + if min_similarity is not None: + search_query.min_similarity = min_similarity + # Presence, not truthiness, for the same reason as everywhere else on this path: + # these are assigned after construction, so the model's own blank guard never + # runs here, and a blank has already been refused by the tool. + if valid_at is not None: + search_query.valid_at = valid_at + if valid_overlaps is not None: + search_query.valid_overlaps = valid_overlaps + if time_kind is not None: + search_query.time_kind = time_kind + + if search_query.no_criteria(): + return None + + # Default to entity-level results to avoid returning individual + # observations/relations as separate search results (see issue #31). + # Applied after no_criteria() so that the implicit default doesn't + # mask a truly empty search request. + if not search_query.entity_types: + # Trigger: a category or valid-time filter was supplied without an + # explicit entity_types. + # Why: both only exist on observations. Categories live on observation + # rows, and temporal assertions are projected against an + # observation's (type, id). Defaulting to "entity" would AND either + # filter against entity rows and return nothing, defeating the + # whole query. + # Outcome: scope the implicit default to observations so + # search_notes(categories=[...]) and search_notes(valid_at=...) + # return the matching bullets. + if search_query.categories or search_query.has_temporal_filter(): + search_query.entity_types = [SearchItemType("observation")] + else: + search_query.entity_types = [SearchItemType("entity")] + return search_query def _compact_search_response(response: SearchResponse) -> SearchResponse: @@ -474,8 +584,28 @@ def _matches_constrained_project(project: dict[str, Any], constrained_project: o return constrained_project in candidates -def _search_project_refs(projects_payload: object) -> list[dict[str, str | None]]: - """Extract project routing refs for optional account-scoped search.""" +@dataclass(frozen=True) +class SearchProjectRef: + """One project an all-projects search may read, and the database it lives in. + + ``name`` is the routable spelling (``workspace/project`` for a cloud project), + which is also the prefix results are qualified with. ``workspace_tenant_id`` is + ``None`` for the local database; every cloud workspace is its own database. + """ + + name: str + external_id: str + id: int + workspace_tenant_id: str | None + path: str + + @property + def bare_name(self) -> str: + return self.name.rsplit("/", 1)[-1] + + +def _search_project_refs(projects_payload: object) -> list[SearchProjectRef]: + """Extract the projects an all-projects search can read from the project list.""" if not isinstance(projects_payload, dict): return [] @@ -484,8 +614,8 @@ def _search_project_refs(projects_payload: object) -> list[dict[str, str | None] if not isinstance(projects, list): return [] - refs: list[dict[str, str | None]] = [] - seen: set[tuple[str | None, str | None]] = set() + refs: list[SearchProjectRef] = [] + seen: set[str] = set() constrained_project = payload.get("constrained_project") for item in projects: if not isinstance(item, dict) or not _matches_constrained_project( @@ -494,42 +624,64 @@ def _search_project_refs(projects_payload: object) -> list[dict[str, str | None] continue project = item.get("qualified_name") or item.get("name") - project_name = project if isinstance(project, str) and project.strip() else None - project_id = _valid_project_id(item.get("external_id")) - if project_name is None and project_id is None: + external_id = _valid_project_id(item.get("external_id")) + internal_id = item.get("id") + # A scoped search addresses a project by its id and attributes hits by its + # external id; a list row missing either cannot be searched this way. + if ( + not isinstance(project, str) + or not project.strip() + or external_id is None + or not isinstance(internal_id, int) + or isinstance(internal_id, bool) + ): continue - - key = (project_name, project_id) - if key in seen: + if external_id in seen: continue - seen.add(key) - refs.append({"project": project_name, "project_id": project_id}) + seen.add(external_id) + tenant = item.get("workspace_tenant_id") + refs.append( + SearchProjectRef( + name=project, + external_id=external_id, + id=internal_id, + workspace_tenant_id=tenant if isinstance(tenant, str) and tenant else None, + path=str(item.get("path") or ""), + ) + ) return refs -async def _load_search_project_refs(context: Context | None = None) -> list[dict[str, str | None]]: +def _select_project_refs( + refs: list[SearchProjectRef], projects: list[str] +) -> list[SearchProjectRef]: + """Keep the projects a caller named, by name, workspace-qualified name, or external id.""" + selected: list[SearchProjectRef] = [] + unknown: list[str] = [] + for requested in projects: + token = requested.strip() + matches = [ref for ref in refs if token in (ref.name, ref.bare_name, ref.external_id)] + if not matches: + unknown.append(token) + continue + for ref in matches: + if ref not in selected: + selected.append(ref) + if unknown: + available = ", ".join(sorted(ref.name for ref in refs)) or "none" + raise ValueError( + f"Unknown project(s): {', '.join(unknown)}. Available projects: {available}" + ) + return selected + + +async def _load_search_project_refs(context: Context | None = None) -> list[SearchProjectRef]: """Load accessible projects for search_all_projects without coupling the wrapper tool.""" from basic_memory.mcp.tools.project_management import list_memory_projects return _search_project_refs(await list_memory_projects(output_format="json", context=context)) -def _raw_results_from_search_payload( - results: SearchResponse | list[SearchResult | dict[str, Any]] | dict[str, Any], -) -> list[SearchResult | dict[str, Any]]: - """Return the result list from any search_notes JSON-compatible payload.""" - if isinstance(results, SearchResponse): - return list(results.results) - if isinstance(results, dict): - nested_results = results.get("results") - return ( - cast(list[SearchResult | dict[str, Any]], nested_results) - if isinstance(nested_results, list) - else [] - ) - return list(results) - - def _result_score(result: SearchResult | dict[str, Any]) -> float: """Return a comparable search score for merged project results.""" if isinstance(result, SearchResult): @@ -561,51 +713,57 @@ def _qualify_permalink_for_project(permalink: object, project: str | None) -> ob ) -def _qualify_results_for_project( - results: list[SearchResult | dict[str, Any]], - project_ref: dict[str, str | None], +def _qualify_result_for_project( + result: SearchResult, + project: str, *, compact: bool = False, +) -> dict[str, Any]: + """Attach the searched workspace/project prefix to one result's permalink.""" + result_data = result.model_dump() + if compact and result_data.get("type") == SearchItemType.OBSERVATION: + # This is an exact file read target, not a generated permalink: + # retain its extension, spaces, and case under the local project or + # workspace/project route. + result_data["permalink"] = f"{project.strip('/')}/{result_data['file_path'].lstrip('/')}" + else: + result_data["permalink"] = _qualify_permalink_for_project( + result_data.get("permalink"), project + ) + return result_data + + +def _attribute_results( + results: list[SearchResult], + refs: list[SearchProjectRef], + *, + compact: bool, ) -> list[dict[str, Any]]: - """Attach the searched workspace/project prefix to each result permalink.""" - qualified: list[dict[str, Any]] = [] + """Qualify each hit's permalink with the project the server says it came from.""" + refs_by_external_id = {ref.external_id: ref for ref in refs} + attributed: list[dict[str, Any]] = [] for result in results: - if isinstance(result, SearchResult): - result_data = result.model_dump() - else: - result_data = dict(result) - project = project_ref.get("project") - if compact and result_data.get("type") == SearchItemType.OBSERVATION and project: - # This is an exact file read target, not a generated permalink: - # retain its extension, spaces, and case under the local project or - # workspace/project route. - result_data["permalink"] = ( - f"{project.strip('/')}/{result_data['file_path'].lstrip('/')}" + ref = refs_by_external_id.get(result.project_external_id or "") + if ref is None: + raise ValueError( + "The search API returned a result it did not attribute to a project in " + "the requested scope. The server is likely older than this client; " + "upgrade it before searching across projects." ) - else: - result_data["permalink"] = _qualify_permalink_for_project( - result_data.get("permalink"), project - ) - qualified.append(result_data) - return qualified - + attributed.append(_qualify_result_for_project(result, ref.name, compact=compact)) + return attributed -def _result_total(results: dict[str, Any], raw_results: list[SearchResult | dict[str, Any]]) -> int: - """Return the best available total for a per-project search payload.""" - total = results.get("total") - if isinstance(total, int) and total > 0: - return total - return len(raw_results) + (1 if results.get("has_more") is True else 0) +def _database_label(tenant_id: str | None, refs: list[SearchProjectRef]) -> str: + """Name one searched database for logs and error responses.""" + if tenant_id is None: + return "local projects" + return f"workspace {refs[0].name.split('/', 1)[0]}" -def _result_total_is_exact(results: dict[str, Any]) -> bool: - """Return whether a per-project payload explicitly guarantees an exact total.""" - return results.get("total_is_exact") is True - -def _project_ref_label(project_ref: dict[str, str | None]) -> str: - """Return a stable log label for a project search ref.""" - return project_ref.get("project") or project_ref.get("project_id") or "" +def _project_item(ref: SearchProjectRef) -> ProjectItem: + """The project shape the readiness check reads (external id and name).""" + return ProjectItem(id=ref.id, external_id=ref.external_id, name=ref.bare_name, path=ref.path) async def _search_all_projects( @@ -628,20 +786,24 @@ async def _search_all_projects( time_kind: str | None, context: Context | None, compact: bool = False, + projects: list[str] | None = None, ) -> dict[str, Any] | str: - """Search every accessible project when the caller explicitly opts in.""" + """Search every accessible project, one query per database. + + Projects that share a database are searched with one scoped query, so within that + database the answer is one ranking rather than per-project pages merged after the + fact. The local database is one group and each cloud workspace is another; only + the merge across databases happens in this process, by score. + """ requested_page = max(page, 1) requested_page_size = max(page_size, 1) - # Each per-project call runs through search_notes -> SearchClient, which refuses a - # response that does not confirm the filter ran. So a project either honored the - # valid-time filter or was dropped with a warning below; the merged answer never - # silently mixes filtered and unfiltered rows. The filter itself is already known to - # be well formed -- search_notes parses it before reaching here -- which is what - # makes "dropped with a warning" mean an unavailable project and nothing else. # Presence, not truthiness: a blank value is refused by `parse_temporal_filter` # before any project is searched, so anything not None is a real question here. temporal_requested = valid_at is not None or valid_overlaps is not None or time_kind is not None project_refs = await _load_search_project_refs(context=context) + if projects: + project_refs = _select_project_refs(project_refs, projects) + scope_label = ", ".join(ref.name for ref in project_refs) if projects else "all projects" if not project_refs: response = SearchResponse( results=[], @@ -654,110 +816,110 @@ async def _search_all_projects( ) if output_format == "json": return response.model_dump(mode="json", exclude_none=True) - return _format_search_markdown(response, "all projects", query) + return _format_search_markdown(response, scope_label, query) + + effective_search_type = search_type or _default_search_type() + search_query = _build_search_query( + query=query, + search_type=effective_search_type, + note_types=note_types, + entity_types=entity_types, + categories=categories, + after_date=after_date, + metadata_filters=metadata_filters, + tags=tags, + status=status, + min_similarity=min_similarity, + valid_at=valid_at, + valid_overlaps=valid_overlaps, + time_kind=time_kind, + ) + if search_query is None: + return _NO_SEARCH_CRITERIA_MESSAGE + query_payload = search_query.model_dump() - per_project_page_size = requested_page * requested_page_size + # Import here to avoid circular import (tools -> clients -> utils -> tools) + from basic_memory.mcp.clients import ScopedSearchClient + + databases: dict[str | None, list[SearchProjectRef]] = {} + for ref in project_refs: + databases.setdefault(ref.workspace_tenant_id, []).append(ref) + + # Each database answers the whole prefix this page needs, so the merge across + # databases can slice it without a second round of requests. + per_database_page_size = requested_page * requested_page_size merged_results: list[dict[str, Any]] = [] total = 0 total_is_exact = True - any_project_has_more = False - # How many projects actually answered. A leg that fails is skipped with a warning, - # so without this the caller cannot tell "no note matched" from "nothing ran". - projects_answered = 0 + any_database_has_more = False + # How many databases actually answered. A database that fails is skipped with a + # warning, so without this the caller cannot tell "no note matched" from + # "nothing ran". + databases_answered = 0 + failures: list[str] = [] query_hint: str | None = None - # Trigger: caller asked for an account-wide search. - # Why: project_id (external UUID) routes through the cloud v2 API path, - # which 401s on local installs because there's no JWT to present. - # Project names route through the local-ASGI path and work for both - # backends — cloud disambiguates names via the workspace/project - # qualified_name already baked into project_ref["project"]. - # Outcome: forward project_id only when the same signals get_project_client - # uses to pick a cloud route are present. Mirrors the cloud_available - # composite in project_context.get_project_client (single source of - # truth for "can we route to cloud?"). - config = ConfigManager().config - use_cloud_routing = ( - is_factory_mode() - or (_explicit_routing() and not _force_local_mode()) - or has_cloud_credentials(config) - ) - - for project_ref in project_refs: - recursive_project_id = project_ref["project_id"] if use_cloud_routing else None + for tenant_id, refs in databases.items(): + label = _database_label(tenant_id, refs) try: - results = await search_notes( - query=query, - project=project_ref["project"], - project_id=recursive_project_id, - page=1, - page_size=per_project_page_size, - search_type=search_type, - output_format="json", - note_types=note_types or None, - entity_types=entity_types or None, - categories=categories or None, - after_date=after_date, - metadata_filters=metadata_filters, - tags=tags, - status=status, - min_similarity=min_similarity, - valid_at=valid_at, - valid_overlaps=valid_overlaps, - time_kind=time_kind, - search_all_projects=False, - context=context, - # Project qualification below must run after compact replaces - # observation excerpt permalinks with their owning file paths. - compact=compact, - ) + async with get_client(workspace=tenant_id) as client: + response = await ScopedSearchClient(client).search( + query_payload, + project_ids=[ref.id for ref in refs], + page=1, + page_size=per_database_page_size, + ) + if not response.results: + # An empty page is a trustworthy miss only after an index pass. + for ref in refs: + guidance = await project_index_required(client, _project_item(ref)) + if guidance is not None: + return guidance except Exception as exc: - logger.warning( - f"Multi-project search failed for project {_project_ref_label(project_ref)}: {exc}" - ) + if _is_service_unavailable_error(exc): + return _format_service_unavailable_response(label, str(exc), query or "") + logger.warning(f"Multi-project search failed for {label}: {exc}") + failures.append(f"{label}: {exc}") total_is_exact = False continue - if isinstance(results, str): - if results.startswith(_SERVICE_UNAVAILABLE_HEADING): - return results - if not results.startswith("# Search Failed"): - return results - logger.warning( - "Multi-project search failed for project " - f"{_project_ref_label(project_ref)}: {results}" - ) - total_is_exact = False - continue - - projects_answered += 1 - if isinstance(results.get("query_hint"), str): - query_hint = results["query_hint"] - raw_results = _raw_results_from_search_payload(results) - total += _result_total(results, raw_results) - total_is_exact = total_is_exact and _result_total_is_exact(results) - any_project_has_more = any_project_has_more or results.get("has_more") is True - merged_results.extend( - _qualify_results_for_project(raw_results, project_ref, compact=compact) + databases_answered += 1 + if compact: + response = _compact_search_response(response) + if response.query_hint: + query_hint = response.query_hint + total += ( + response.total + if response.total > 0 + else len(response.results) + (1 if response.has_more else 0) ) - - # Trigger: a valid-time filter was requested and not one project answered. - # Why: each leg confirms the filter through SearchClient or is refused by it, and a - # refusal is caught above, logged, and skipped -- so a fleet of servers predating - # SPEC-82 drops every leg and arrives here indistinguishable from "no note matched". - # Claiming `temporal_applied` on that would confirm a filter that ran nowhere, which - # is the version skew the client's own check exists to make loud. + total_is_exact = total_is_exact and response.total_is_exact + any_database_has_more = any_database_has_more or response.has_more + merged_results.extend(_attribute_results(response.results, refs, compact=compact)) + + # Trigger: a valid-time filter was requested and not one database answered. + # Why: each database confirms the filter through the client or is refused by it, + # and a refusal is caught above, logged, and skipped -- so a fleet of servers + # predating SPEC-82 drops every database and arrives here indistinguishable from + # "no note matched". Claiming `temporal_applied` on that would confirm a filter + # that ran nowhere, which is the version skew the client's own check exists to + # make loud. # Outcome: the skew is propagated as one error naming it, rather than returning an # empty result wearing the shape of a successful filtered search. - if temporal_requested and projects_answered == 0: + if temporal_requested and databases_answered == 0: raise ValueError( "No project applied the requested valid-time filter: every project was " "skipped, so the filter ran nowhere and an empty result would not mean " "'no matches'. The servers are likely older than this client; upgrade them " "or drop valid_at / valid_overlaps / time_kind from the query." ) + if databases_answered == 0: + # Every database failed. That is a failed search, not an empty one. + return _format_search_error_response( + scope_label, "; ".join(failures), query or "", effective_search_type + ) - # Each project owns retrieval and optional reranking behind its typed API client. + # Each database owns retrieval and optional reranking behind its typed API client. # The MCP process only merges returned scores; it must not instantiate repository # providers with local credentials for content fetched through another route. sorted_results = sorted(merged_results, key=_result_score, reverse=True) @@ -767,17 +929,17 @@ async def _search_all_projects( response = SearchResponse.model_validate( { "results": paged_results, - # Only propagate query guidance when every project answered and the - # aggregate is empty; one empty leg must not label a successful search. + # Only propagate query guidance when every database answered and the + # aggregate is empty; one empty database must not label a successful search. "query_hint": query_hint - if requested_page == 1 and not merged_results and projects_answered == len(project_refs) + if requested_page == 1 and not merged_results and databases_answered == len(databases) else None, "current_page": requested_page, "page_size": requested_page_size, "total": total, "total_is_exact": total_is_exact, - "has_more": any_project_has_more or total > end or len(sorted_results) > end, - # Confirmed only because a project answered: `projects_answered` is + "has_more": any_database_has_more or total > end or len(sorted_results) > end, + # Confirmed only because a database answered: `databases_answered` is # non-zero here for any temporal query, guarded immediately above. "temporal_applied": True if temporal_requested else None, } @@ -785,7 +947,7 @@ async def _search_all_projects( if output_format == "json": return response.model_dump(mode="json", exclude_none=True) - return _format_search_markdown(response, "all projects", query) + return _format_search_markdown(response, scope_label, query) @mcp.tool( @@ -817,6 +979,7 @@ async def search_notes( validation_alias=AliasChoices("search_all_projects", "all_projects"), ), ] = False, + projects: Optional[List[str]] = None, # `offset` is intentionally NOT aliased to `page`: offset is item-indexed # (skip N items) while page is 1-indexed page-number. Direct aliasing would # silently return the wrong slice. @@ -941,8 +1104,8 @@ async def search_notes( Project Resolution: Server resolves projects in this order: Single Project Mode → project parameter → default project. If project unknown, use list_memory_projects() or recent_activity() first. - Set search_all_projects=True to search every accessible project; this is opt-in because it - performs one search per project. + Set search_all_projects=True to search every accessible project, or pass projects=[...] to + search a chosen subset. Either runs one scoped query per database the projects live in. ## Search Syntax Examples @@ -1083,6 +1246,10 @@ async def search_notes( workspaces. Takes precedence over `project`. Get from list_memory_projects(). search_all_projects: Optional opt-in to search every accessible project. Ignored when `project` or `project_id` is supplied. + projects: Optional list of project names or external ids to search together. Names + are matched exactly as list_memory_projects() reports them (cloud projects + by their workspace-qualified name). Ignored when `project` or `project_id` + is supplied; an unknown name is an error rather than a silent skip. page: The page number of results to return (default 1). Aliases: page_number. page_size: The number of results to return per page (default 10). @@ -1236,6 +1403,7 @@ async def search_notes( # Outcome: comma-split/list normalization applies on every path; parse_str_list is # idempotent, so MCP-validated input passes through unchanged. note_types = parse_str_list(note_types) if note_types is not None else [] + projects = parse_str_list(projects) if projects is not None else None entity_types = parse_str_list(entity_types) if entity_types is not None else [] categories = parse_str_list(categories) if categories is not None else [] @@ -1290,7 +1458,7 @@ async def search_notes( # already provide a concrete project route. # Why: multi-project fan-out can be slow, so default search remains project-scoped. # Outcome: run one normal search per accessible project and merge ranked results. - if search_all_projects and project is None and project_id is None: + if (search_all_projects or projects) and project is None and project_id is None: all_projects_result = await _search_all_projects( query=query, page=page, @@ -1310,6 +1478,7 @@ async def search_notes( time_kind=time_kind, context=context, compact=compact, + projects=projects, ) return all_projects_result @@ -1365,104 +1534,24 @@ async def search_notes( effective_search_type = "permalink" try: - # Create a SearchQuery object based on the parameters - search_query = SearchQuery() - - # Only map search_type to query fields when there is an actual query string. - # When query is None/empty, skip the search mode block — filters-only path. + search_query = _build_search_query( + query=query, + search_type=effective_search_type, + note_types=note_types, + entity_types=entity_types, + categories=categories, + after_date=after_date, + metadata_filters=metadata_filters, + tags=tags, + status=status, + min_similarity=min_similarity, + valid_at=valid_at, + valid_overlaps=valid_overlaps, + time_kind=time_kind, + ) + if search_query is None: + return _NO_SEARCH_CRITERIA_MESSAGE effective_query = (query or "").strip() - if effective_query: - valid_search_types = { - "text", - "title", - "permalink", - "vector", - "semantic", - "hybrid", - } - if effective_search_type == "text": - search_query.text = effective_query - search_query.retrieval_mode = SearchRetrievalMode.FTS - elif effective_search_type in ("vector", "semantic"): - search_query.text = effective_query - search_query.retrieval_mode = SearchRetrievalMode.VECTOR - elif effective_search_type == "hybrid": - search_query.text = effective_query - search_query.retrieval_mode = SearchRetrievalMode.HYBRID - elif effective_search_type == "title": - search_query.title = effective_query - elif effective_search_type == "permalink" and "*" in effective_query: - search_query.permalink_match = effective_query - elif effective_search_type == "permalink": - search_query.permalink = effective_query - else: - raise ValueError( - f"Invalid search_type '{effective_search_type}'. " - f"Valid options: {', '.join(sorted(valid_search_types))}" - ) - - # Add optional filters if provided (empty lists are treated as no filter) - if entity_types: - search_query.entity_types = [SearchItemType(t) for t in entity_types] - if categories: - search_query.categories = categories - if note_types: - search_query.note_types = note_types - if after_date: - search_query.after_date = after_date - if metadata_filters: - # Alias common column/model names to their frontmatter key equivalents. - # Users often pass "note_type" (the entity model column) when the - # frontmatter field is actually "type". - _METADATA_KEY_ALIASES = {"note_type": "type"} - metadata_filters = { - _METADATA_KEY_ALIASES.get(k, k): v for k, v in metadata_filters.items() - } - search_query.metadata_filters = metadata_filters - if tags: - search_query.tags = tags - if status: - search_query.status = status - if min_similarity is not None: - search_query.min_similarity = min_similarity - # Presence, not truthiness, for the same reason as everywhere else on - # this path: these are assigned after construction, so the model's own - # blank guard never runs here, and a blank has already been refused above. - if valid_at is not None: - search_query.valid_at = valid_at - if valid_overlaps is not None: - search_query.valid_overlaps = valid_overlaps - if time_kind is not None: - search_query.time_kind = time_kind - - # Reject searches with no criteria at all - if search_query.no_criteria(): - return ( - "# No Search Criteria\n\n" - "Please provide at least one of: `query`, `metadata_filters`, " - "`tags`, `status`, `note_types`, `entity_types`, `categories`, " - "`after_date`, `valid_at`, `valid_overlaps`, or `time_kind`." - ) - - # Default to entity-level results to avoid returning individual - # observations/relations as separate search results (see issue #31). - # Applied after no_criteria() so that the implicit default doesn't - # mask a truly empty search request. - if not search_query.entity_types: - # Trigger: a category or valid-time filter was supplied without an - # explicit entity_types. - # Why: both only exist on observations — categories live on observation - # rows, and temporal assertions are projected against an - # observation's (type, id). Defaulting to "entity" would AND either - # filter against entity rows and return nothing, defeating the - # whole query. - # Outcome: scope the implicit default to observations so - # search_notes(categories=[...]) and search_notes(valid_at=...) - # return the matching bullets. - if search_query.categories or search_query.has_temporal_filter(): - search_query.entity_types = [SearchItemType("observation")] - else: - search_query.entity_types = [SearchItemType("entity")] logger.debug( f"Search request: project={active_project.name} " diff --git a/tests/mcp/test_tool_contracts.py b/tests/mcp/test_tool_contracts.py index 43b828995..f6693eae3 100644 --- a/tests/mcp/test_tool_contracts.py +++ b/tests/mcp/test_tool_contracts.py @@ -136,6 +136,7 @@ "project", "project_id", "search_all_projects", + "projects", "page", "page_size", "search_type", diff --git a/tests/mcp/test_tool_search_temporal.py b/tests/mcp/test_tool_search_temporal.py index c5a5b388d..45a380adb 100644 --- a/tests/mcp/test_tool_search_temporal.py +++ b/tests/mcp/test_tool_search_temporal.py @@ -11,12 +11,14 @@ """ import inspect +import sys from typing import Any +from unittest.mock import AsyncMock import pytest from basic_memory.mcp.tools import write_note -from basic_memory.mcp.tools.search import search_notes +from basic_memory.mcp.tools.search import SearchProjectRef, search_notes # The spec's worked example, verbatim: one note, two decisions, adjacent half-open # effective windows meeting at the July 27 cutover. @@ -474,17 +476,37 @@ async def test_all_projects_search_propagates_a_filter_no_project_could_apply( """ await _write_cache_layer_note(test_project.name) - import sys - # `basic_memory.mcp.tools.search` is shadowed by a `search` function exported # from the package, so reach the module itself rather than that name. search_module = sys.modules["basic_memory.mcp.tools.search"] - async def refuse_every_leg(*args: Any, **kwargs: Any) -> str: - # What a per-project leg looks like once SearchClient rejects the response. - return "# Search Failed\n\nThe search API did not apply the requested valid-time filter" + monkeypatch.setattr( + search_module, + "_load_search_project_refs", + AsyncMock( + return_value=[ + SearchProjectRef( + name=test_project.name, + external_id=test_project.external_id, + id=test_project.id, + workspace_tenant_id=None, + path=test_project.path, + ) + ] + ), + ) + + class RefusingScopedSearchClient: + # What a database looks like once the client rejects its response. + def __init__(self, client: Any) -> None: + pass + + async def search(self, *args: Any, **kwargs: Any) -> Any: + raise ValueError("The search API did not apply the requested valid-time filter") - monkeypatch.setattr(search_module, "search_notes", refuse_every_leg) + monkeypatch.setattr( + sys.modules["basic_memory.mcp.clients"], "ScopedSearchClient", RefusingScopedSearchClient + ) with pytest.raises(ValueError, match="No project applied the requested valid-time filter"): await search_module._search_all_projects( diff --git a/tests/mcp/tools/test_cjk_search_guidance.py b/tests/mcp/tools/test_cjk_search_guidance.py index 6507fd10b..3da88ee67 100644 --- a/tests/mcp/tools/test_cjk_search_guidance.py +++ b/tests/mcp/tools/test_cjk_search_guidance.py @@ -1,51 +1,87 @@ """Account search carries query guidance only for a complete, empty first page.""" import importlib +from contextlib import asynccontextmanager from typing import Literal from unittest.mock import AsyncMock import pytest +from basic_memory.mcp.tools.search import SearchProjectRef +from basic_memory.schemas.search import SearchItemType, SearchResponse, SearchResult + +ONE = SearchProjectRef( + name="one", + external_id="11111111-1111-1111-1111-111111111111", + id=1, + workspace_tenant_id=None, + path="/one", +) +TWO = SearchProjectRef( + name="two", + external_id="22222222-2222-2222-2222-222222222222", + id=2, + workspace_tenant_id="tenant-two", + path="/two", +) + @pytest.mark.asyncio @pytest.mark.parametrize("case", ["empty", "hit", "failed", "later"]) @pytest.mark.parametrize("output_format", ["text", "json"]) async def test_fanout_query_hint(monkeypatch, case: str, output_format: Literal["text", "json"]): search = importlib.import_module("basic_memory.mcp.tools.search") - monkeypatch.setattr( - search, - "_load_search_project_refs", - AsyncMock( - return_value=[ - {"project": "one", "project_id": None}, - {"project": "two", "project_id": None}, - ] - ), - ) - empty = { - "results": [], - "total": 0, - "total_is_exact": True, - "query_hint": "Try a shorter word from your query.", - } - second: dict[str, object] | str = empty - if case == "hit": - second = { - "results": [ - { - "title": "Match", - "type": "entity", - "score": 1.0, - "file_path": "note.md", - "permalink": "note", - } - ], - "total": 1, - "total_is_exact": True, - } - elif case == "failed": - second = "# Search Failed - Access Error" - monkeypatch.setattr(search, "search_notes", AsyncMock(side_effect=[empty, second])) + clients = importlib.import_module("basic_memory.mcp.clients") + monkeypatch.setattr(search, "_load_search_project_refs", AsyncMock(return_value=[ONE, TWO])) + monkeypatch.setattr(search, "project_index_required", AsyncMock(return_value=None)) + + @asynccontextmanager + async def fake_get_client(project_name=None, workspace=None): + yield {"workspace": workspace} + + def empty(page: int, page_size: int) -> SearchResponse: + return SearchResponse( + results=[], + current_page=page, + page_size=page_size, + total=0, + total_is_exact=True, + query_hint="Try a shorter word from your query.", + ) + + class StubScopedSearchClient: + def __init__(self, client): + self.workspace = client["workspace"] + + async def search(self, payload, *, project_ids, page, page_size): + # The first database (local) is always empty; the second database varies. + if self.workspace is None: + return empty(page, page_size) + if case == "hit": + return SearchResponse( + results=[ + SearchResult( + title="Match", + type=SearchItemType.ENTITY, + score=1.0, + file_path="note.md", + permalink="note", + project_id=TWO.id, + project_external_id=TWO.external_id, + ) + ], + current_page=page, + page_size=page_size, + total=1, + total_is_exact=True, + ) + if case == "failed": + raise RuntimeError("access denied") + return empty(page, page_size) + + monkeypatch.setattr(search, "get_client", fake_get_client) + monkeypatch.setattr(clients, "ScopedSearchClient", StubScopedSearchClient) + result = await search._search_all_projects( query="雾凇拼音", page=2 if case == "later" else 1, diff --git a/tests/mcp/tools/test_search_notes_multi_project.py b/tests/mcp/tools/test_search_notes_multi_project.py index 69a616cf0..27541c55c 100644 --- a/tests/mcp/tools/test_search_notes_multi_project.py +++ b/tests/mcp/tools/test_search_notes_multi_project.py @@ -1,116 +1,158 @@ -"""Tests for optional multi-project search_notes behavior.""" +"""Optional multi-project search_notes behavior: one scoped query per database.""" -from contextlib import asynccontextmanager import importlib +from contextlib import asynccontextmanager +from typing import Any from unittest.mock import AsyncMock -from httpx import HTTPStatusError, Request, Response -from fastmcp.exceptions import ToolError import pytest +from fastmcp.exceptions import ToolError +from httpx import HTTPStatusError, Request, Response +from basic_memory.mcp.tools.search import SearchProjectRef from basic_memory.schemas.search import SearchItemType, SearchResponse, SearchResult +PERSONAL = SearchProjectRef( + name="personal/main", + external_id="11111111-1111-1111-1111-111111111111", + id=11, + workspace_tenant_id="tenant-personal", + path="/personal/main", +) +TEAM = SearchProjectRef( + name="team-paul/main", + external_id="22222222-2222-2222-2222-222222222222", + id=22, + workspace_tenant_id="tenant-team", + path="/team/main", +) +ALPHA = SearchProjectRef( + name="alpha", + external_id="33333333-3333-3333-3333-333333333333", + id=3, + workspace_tenant_id=None, + path="/alpha", +) +BETA = SearchProjectRef( + name="beta", + external_id="44444444-4444-4444-4444-444444444444", + id=4, + workspace_tenant_id=None, + path="/beta", +) + + +def _result( + ref: SearchProjectRef, + *, + title: str, + score: float, + observation: bool = False, + permalink: str = "notes/example", + file_path: str = "/notes/example.md", +) -> SearchResult: + return SearchResult( + title=title, + permalink=permalink, + content="MCP content", + type=SearchItemType.OBSERVATION if observation else SearchItemType.ENTITY, + category="fact" if observation else None, + score=score, + file_path=file_path, + project_id=ref.id, + project_external_id=ref.external_id, + ) + -def _stub_routing_mode(monkeypatch, *, cloud: bool) -> None: - """Pin the three cloud-route signals search.py reads. +def _install_scoped_search( + monkeypatch: pytest.MonkeyPatch, + refs: list[SearchProjectRef], + answer, +) -> tuple[list[str | None], list[dict[str, Any]]]: + """Route the all-projects search into a stub that records each database's call. - `_search_all_projects` only forwards project_id (external UUID) when a - cloud route is available. The composite mirrors get_project_client: - factory mode OR explicit --cloud OR has_cloud_credentials. Tests stub - all three so a dev box with OAuth tokens on disk can't bleed into the - local-mode case. + ``answer(workspace, project_ids, page, page_size)`` returns the SearchResponse for + one database, or raises to model a failed database. """ + clients_mod = importlib.import_module("basic_memory.mcp.clients") search_mod = importlib.import_module("basic_memory.mcp.tools.search") - monkeypatch.setattr(search_mod, "is_factory_mode", lambda: False) - monkeypatch.setattr(search_mod, "_explicit_routing", lambda: cloud) - monkeypatch.setattr(search_mod, "_force_local_mode", lambda: False) - monkeypatch.setattr(search_mod, "has_cloud_credentials", lambda config: cloud) + workspaces: list[str | None] = [] + calls: list[dict[str, Any]] = [] + async def fake_load_search_project_refs(context=None): + return refs -@pytest.fixture -def cloud_routing(monkeypatch): - """Force the cloud-routing path for multi-project search tests.""" - _stub_routing_mode(monkeypatch, cloud=True) + @asynccontextmanager + async def fake_get_client(project_name=None, workspace=None): + workspaces.append(workspace) + yield {"workspace": workspace} + + class StubScopedSearchClient: + def __init__(self, client): + self.workspace = client["workspace"] + + async def search(self, payload, *, project_ids, page, page_size): + calls.append( + { + "workspace": self.workspace, + "project_ids": list(project_ids), + "page": page, + "page_size": page_size, + "payload": payload, + } + ) + return answer(self.workspace, list(project_ids), page, page_size) + + monkeypatch.setattr(search_mod, "_load_search_project_refs", fake_load_search_project_refs) + monkeypatch.setattr(search_mod, "get_client", fake_get_client) + monkeypatch.setattr(clients_mod, "ScopedSearchClient", StubScopedSearchClient) + monkeypatch.setattr(search_mod, "project_index_required", AsyncMock(return_value=None)) + return workspaces, calls -@pytest.fixture -def local_routing(monkeypatch): - """Force the local-routing path for multi-project search tests.""" - _stub_routing_mode(monkeypatch, cloud=False) +def _page(results: list[SearchResult], page: int, page_size: int, **fields: Any) -> SearchResponse: + return SearchResponse( + results=results, + current_page=page, + page_size=page_size, + total=fields.pop("total", len(results)), + **fields, + ) @pytest.mark.asyncio @pytest.mark.parametrize("compact", [False, True]) @pytest.mark.parametrize("observations", [False, True]) -async def test_search_notes_search_all_projects_qualifies_result_permalinks( - monkeypatch, cloud_routing, compact, observations +async def test_each_workspace_is_one_query_and_results_stay_routable( + monkeypatch, compact, observations ): - """Multi-project search belongs to search_notes and keeps result ids routable.""" - clients_mod = importlib.import_module("basic_memory.mcp.clients") + """Two cloud workspaces are two databases: one scoped call each, merged by score.""" search_mod = importlib.import_module("basic_memory.mcp.tools.search") - project_refs = [ - { - "project": "personal/main", - "project_id": "11111111-1111-1111-1111-111111111111", - }, - { - "project": "team-paul/main", - "project_id": "22222222-2222-2222-2222-222222222222", - }, - ] - searched_projects: list[tuple[str | None, str | None]] = [] - - async def fake_load_search_project_refs(context=None): - return project_refs - - class StubProject: - def __init__(self, name: str | None, external_id: str | None): - self.name = name or "main" - self.external_id = external_id or "local-main" - - @asynccontextmanager - async def fake_get_project_client(project=None, context=None, project_id=None): - searched_projects.append((project, project_id)) - yield object(), StubProject(project, project_id) - - async def fake_resolve_project_and_path(client, identifier, project=None, context=None): - return StubProject(project, None), identifier, False - - class MockSearchClient: - def __init__(self, client, project_id): - self.project_id = project_id + def answer(workspace, project_ids, page, page_size): + ref, title, score = ( + (PERSONAL, "Personal MCP Test Note", 0.5) + if workspace == PERSONAL.workspace_tenant_id + else (TEAM, "Team MCP Test Note", 0.9) + ) + return _page( + [ + _result( + ref, + title=title, + score=score, + observation=observations, + permalink="main/tests/mcp-test-note", + file_path="tests/Exact Note.md" + if observations + else "/main/tests/mcp-test-note.md", + ) + ], + page, + page_size, + ) - async def search(self, payload, page, page_size): - if self.project_id == "11111111-1111-1111-1111-111111111111": - title = "Personal MCP Test Note" - score = 0.5 - else: - title = "Team MCP Test Note" - score = 0.9 - return SearchResponse( - results=[ - SearchResult( - title=title, - permalink="main/tests/mcp-test-note", - content="MCP content", - type=SearchItemType.OBSERVATION if observations else SearchItemType.ENTITY, - category="fact" if observations else None, - score=score, - file_path="tests/Exact Note.md" - if observations - else "/main/tests/mcp-test-note.md", - ) - ], - current_page=page, - page_size=page_size, - total=1, - ) - - monkeypatch.setattr(search_mod, "_load_search_project_refs", fake_load_search_project_refs) - monkeypatch.setattr(search_mod, "get_project_client", fake_get_project_client) - monkeypatch.setattr(search_mod, "resolve_project_and_path", fake_resolve_project_and_path) - monkeypatch.setattr(clients_mod, "SearchClient", MockSearchClient) + workspaces, calls = _install_scoped_search(monkeypatch, [PERSONAL, TEAM], answer) result = await search_mod.search_notes( query="MCP Test Note", @@ -120,9 +162,10 @@ async def search(self, payload, page, page_size): ) assert isinstance(result, dict) - assert searched_projects == [ - ("personal/main", "11111111-1111-1111-1111-111111111111"), - ("team-paul/main", "22222222-2222-2222-2222-222222222222"), + assert workspaces == [PERSONAL.workspace_tenant_id, TEAM.workspace_tenant_id] + assert [(call["workspace"], call["project_ids"]) for call in calls] == [ + (PERSONAL.workspace_tenant_id, [PERSONAL.id]), + (TEAM.workspace_tenant_id, [TEAM.id]), ] path = "tests/Exact Note.md" if compact and observations else "tests/mcp-test-note" assert [item["permalink"] for item in result["results"]] == [ @@ -131,11 +174,101 @@ async def search(self, payload, page, page_size): ] if compact: assert all("content" not in item for item in result["results"]) + assert result["total"] == 2 + assert result["total_is_exact"] is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("compact_observations", [False, True]) +async def test_local_projects_share_one_scoped_query(monkeypatch, compact_observations): + """Projects in the local database are one query, and every hit names its project.""" + search_mod = importlib.import_module("basic_memory.mcp.tools.search") + + def answer(workspace, project_ids, page, page_size): + return _page( + [ + _result( + ref, title=f"Note in {ref.name}", score=0.5, observation=compact_observations + ) + for ref in (ALPHA, BETA) + ], + page, + page_size, + ) + + workspaces, calls = _install_scoped_search(monkeypatch, [ALPHA, BETA], answer) + + result = await search_mod.search_notes( + query="anything", + search_all_projects=True, + output_format="json", + compact=compact_observations, + ) + + assert isinstance(result, dict) + assert workspaces == [None] + assert [(call["workspace"], call["project_ids"]) for call in calls] == [ + (None, [ALPHA.id, BETA.id]) + ] + assert result["total"] == 2 assert result["total_is_exact"] is True + if compact_observations: + assert [item["permalink"] for item in result["results"]] == [ + "alpha/notes/example.md", + "beta/notes/example.md", + ] + else: + assert [item["permalink"] for item in result["results"]] == [ + "notes/example", + "notes/example", + ] + assert [item["project_external_id"] for item in result["results"]] == [ + ALPHA.external_id, + BETA.external_id, + ] + + +@pytest.mark.asyncio +async def test_each_database_answers_the_whole_prefix_the_page_needs(monkeypatch): + """Page two of five asks every database for its top ten, then slices the merge.""" + search_mod = importlib.import_module("basic_memory.mcp.tools.search") + + def answer(workspace, project_ids, page, page_size): + ref = PERSONAL if workspace == PERSONAL.workspace_tenant_id else TEAM + # Distinct scores across databases so the merge order is unambiguous. + base = 0.9 if ref is TEAM else 0.85 + return _page( + [ + _result(ref, title=f"{ref.name} {index}", score=base - index * 0.1) + for index in range(6) + ], + page, + page_size, + total=6, + ) + + _workspaces, calls = _install_scoped_search(monkeypatch, [PERSONAL, TEAM], answer) + + result = await search_mod.search_notes( + query="notes", search_all_projects=True, output_format="json", page=2, page_size=5 + ) + + assert isinstance(result, dict) + assert [(call["page"], call["page_size"]) for call in calls] == [(1, 10), (1, 10)] + assert [item["title"] for item in result["results"]] == [ + "personal/main 2", + "team-paul/main 3", + "personal/main 3", + "team-paul/main 4", + "personal/main 4", + ] + assert result["current_page"] == 2 + assert result["total"] == 12 + assert result["has_more"] is True @pytest.mark.asyncio -async def test_search_notes_multi_project_search_is_opt_in(monkeypatch): +async def test_multi_project_search_is_opt_in(monkeypatch): """Default search_notes calls stay scoped to the resolved project.""" clients_mod = importlib.import_module("basic_memory.mcp.clients") search_mod = importlib.import_module("basic_memory.mcp.tools.search") @@ -168,7 +301,6 @@ async def search(self, payload, page, page_size): monkeypatch.setattr(search_mod, "get_project_client", fake_get_project_client) monkeypatch.setattr(search_mod, "resolve_project_and_path", fake_resolve_project_and_path) monkeypatch.setattr(clients_mod, "SearchClient", MockSearchClient) - monkeypatch.setattr(search_mod, "project_index_required", AsyncMock(return_value=None)) result = await search_mod.search_notes(query="MCP Test Note", output_format="json") @@ -179,9 +311,7 @@ async def search(self, payload, page, page_size): @pytest.mark.asyncio -async def test_search_notes_search_all_projects_with_no_refs_returns_empty_all_projects( - monkeypatch, -): +async def test_no_accessible_projects_is_an_empty_all_projects_answer(monkeypatch): """Explicit all-project search must not silently fall back to one project.""" search_mod = importlib.import_module("basic_memory.mcp.tools.search") @@ -189,12 +319,12 @@ async def fake_load_search_project_refs(context=None): return [] @asynccontextmanager - async def fail_get_project_client(*args, **kwargs): - raise AssertionError("search_all_projects=True should not fall back to scoped search") + async def fail_get_client(*args, **kwargs): + raise AssertionError("search_all_projects=True should not query without projects") yield monkeypatch.setattr(search_mod, "_load_search_project_refs", fake_load_search_project_refs) - monkeypatch.setattr(search_mod, "get_project_client", fail_get_project_client) + monkeypatch.setattr(search_mod, "get_client", fail_get_client) result = await search_mod.search_notes( query="MCP Test Note", @@ -202,7 +332,6 @@ async def fail_get_project_client(*args, **kwargs): output_format="json", ) - assert isinstance(result, dict) assert result == { "results": [], "current_page": 1, @@ -214,40 +343,11 @@ async def fail_get_project_client(*args, **kwargs): @pytest.mark.asyncio -async def test_search_notes_search_all_projects_continues_after_project_failure( - monkeypatch, cloud_routing -): - """One failing project should not discard successful all-project search results.""" - clients_mod = importlib.import_module("basic_memory.mcp.clients") +async def test_a_failing_database_is_skipped_and_the_total_becomes_inexact(monkeypatch): + """One failing workspace should not discard the other's results.""" search_mod = importlib.import_module("basic_memory.mcp.tools.search") - - project_refs = [ - { - "project": "personal/main", - "project_id": "11111111-1111-1111-1111-111111111111", - }, - { - "project": "team-paul/main", - "project_id": "22222222-2222-2222-2222-222222222222", - }, - ] warnings: list[str] = [] - async def fake_load_search_project_refs(context=None): - return project_refs - - class StubProject: - def __init__(self, name: str | None, external_id: str | None): - self.name = name or "main" - self.external_id = external_id or "local-main" - - @asynccontextmanager - async def fake_get_project_client(project=None, context=None, project_id=None): - yield object(), StubProject(project, project_id) - - async def fake_resolve_project_and_path(client, identifier, project=None, context=None): - return StubProject(project, None), identifier, False - class FakeLogger: def debug(self, *args, **kwargs): pass @@ -258,34 +358,15 @@ def error(self, *args, **kwargs): def warning(self, message, *args, **kwargs): warnings.append(str(message)) - class MockSearchClient: - def __init__(self, client, project_id): - self.project_id = project_id - - async def search(self, payload, page, page_size): - if self.project_id == "22222222-2222-2222-2222-222222222222": - raise RuntimeError("team index unavailable") - return SearchResponse( - results=[ - SearchResult( - title="Personal MCP Test Note", - permalink="main/tests/mcp-test-note", - content="MCP content", - type=SearchItemType.ENTITY, - score=0.5, - file_path="/main/tests/mcp-test-note.md", - ) - ], - current_page=page, - page_size=page_size, - total=1, - ) + def answer(workspace, project_ids, page, page_size): + if workspace == TEAM.workspace_tenant_id: + raise RuntimeError("team index unavailable") + return _page( + [_result(PERSONAL, title="Personal MCP Test Note", score=0.5)], page, page_size + ) - monkeypatch.setattr(search_mod, "_load_search_project_refs", fake_load_search_project_refs) - monkeypatch.setattr(search_mod, "get_project_client", fake_get_project_client) - monkeypatch.setattr(search_mod, "resolve_project_and_path", fake_resolve_project_and_path) + _install_scoped_search(monkeypatch, [PERSONAL, TEAM], answer) monkeypatch.setattr(search_mod, "logger", FakeLogger()) - monkeypatch.setattr(clients_mod, "SearchClient", MockSearchClient) result = await search_mod.search_notes( query="MCP Test Note", @@ -294,84 +375,50 @@ async def search(self, payload, page, page_size): ) assert isinstance(result, dict) - assert [item["permalink"] for item in result["results"]] == [ - "personal/main/tests/mcp-test-note", - ] + assert [item["permalink"] for item in result["results"]] == ["personal/main/notes/example"] assert result["total"] == 1 assert result["total_is_exact"] is False - assert any("team-paul/main" in warning for warning in warnings) + assert any("workspace team-paul" in warning for warning in warnings) assert any("team index unavailable" in warning for warning in warnings) @pytest.mark.asyncio -async def test_search_notes_search_all_projects_propagates_retryable_service_outage( - monkeypatch, cloud_routing -): - """A retryable project outage must fail the merged page instead of returning a partial one.""" - clients_mod = importlib.import_module("basic_memory.mcp.clients") +async def test_every_database_failing_is_a_failed_search_not_an_empty_one(monkeypatch): + """With one database, its failure is the whole answer; do not dress it as no matches.""" search_mod = importlib.import_module("basic_memory.mcp.tools.search") - project_refs = [ - { - "project": "personal/main", - "project_id": "11111111-1111-1111-1111-111111111111", - }, - { - "project": "team-paul/main", - "project_id": "22222222-2222-2222-2222-222222222222", - }, - ] - async def fake_load_search_project_refs(context=None): - return project_refs + def answer(workspace, project_ids, page, page_size): + raise RuntimeError("local index unavailable") - class StubProject: - def __init__(self, name: str | None, external_id: str | None): - self.name = name or "main" - self.external_id = external_id or "local-main" + _install_scoped_search(monkeypatch, [ALPHA, BETA], answer) - @asynccontextmanager - async def fake_get_project_client(project=None, context=None, project_id=None): - yield object(), StubProject(project, project_id) + result = await search_mod.search_notes( + query="MCP Test Note", search_all_projects=True, output_format="json" + ) - async def fake_resolve_project_and_path(client, identifier, project=None, context=None): - return StubProject(project, None), identifier, False + assert isinstance(result, str) + assert result.startswith("# Search Failed") + assert "local index unavailable" in result - class MockSearchClient: - def __init__(self, client, project_id): - self.project_id = project_id - async def search(self, payload, page, page_size): - if self.project_id == "22222222-2222-2222-2222-222222222222": - request = Request("POST", "https://api.example/search") - response = Response( - 503, - request=request, - json={"detail": "Reranker temporarily unavailable"}, - ) - try: - response.raise_for_status() - except HTTPStatusError as exc: - raise ToolError("Reranker temporarily unavailable") from exc - return SearchResponse( - results=[ - SearchResult( - title="Personal result", - permalink="main/personal-result", - content="MCP content", - type=SearchItemType.ENTITY, - score=0.5, - file_path="/main/personal-result.md", - ) - ], - current_page=page, - page_size=page_size, - total=1, +@pytest.mark.asyncio +async def test_a_retryable_service_outage_fails_the_merged_page(monkeypatch): + """A retryable outage must fail the merged page instead of returning a partial one.""" + search_mod = importlib.import_module("basic_memory.mcp.tools.search") + + def answer(workspace, project_ids, page, page_size): + if workspace == TEAM.workspace_tenant_id: + request = Request("QUERY", "https://api.example/v2/search/") + response = Response( + 503, request=request, json={"detail": "Reranker temporarily unavailable"} ) + try: + response.raise_for_status() + except HTTPStatusError as exc: + raise ToolError("Reranker temporarily unavailable") from exc + return _page([_result(PERSONAL, title="Personal result", score=0.5)], page, page_size) - monkeypatch.setattr(search_mod, "_load_search_project_refs", fake_load_search_project_refs) - monkeypatch.setattr(search_mod, "get_project_client", fake_get_project_client) - monkeypatch.setattr(search_mod, "resolve_project_and_path", fake_resolve_project_and_path) - monkeypatch.setattr(clients_mod, "SearchClient", MockSearchClient) + _install_scoped_search(monkeypatch, [PERSONAL, TEAM], answer) result = await search_mod.search_notes( query="MCP Test Note", @@ -384,92 +431,118 @@ async def search(self, payload, page, page_size): @pytest.mark.asyncio -@pytest.mark.parametrize("compact_observations", [False, True]) -async def test_search_notes_search_all_projects_local_omits_project_id( - monkeypatch, local_routing, compact_observations -): - """Without a cloud route, fan-out must address each project by name only. - - project_id (external UUID) routes through the cloud v2 API path, which - returns 401 on local installs because there's no JWT to present. Local - fan-out has to fall back to the name-routed path so each per-project - search actually returns results instead of silently failing. - """ - clients_mod = importlib.import_module("basic_memory.mcp.clients") +async def test_an_empty_answer_checks_every_project_in_the_database_for_an_index(monkeypatch): + """An empty scoped page is only a miss once each project has been indexed.""" search_mod = importlib.import_module("basic_memory.mcp.tools.search") - project_refs = [ - { - "project": "alpha", - "project_id": "11111111-1111-1111-1111-111111111111", - }, - { - "project": "beta", - "project_id": "22222222-2222-2222-2222-222222222222", - }, - ] - searched_projects: list[tuple[str | None, str | None]] = [] + def answer(workspace, project_ids, page, page_size): + return _page([], page, page_size, total=0) - async def fake_load_search_project_refs(context=None): - return project_refs + _install_scoped_search(monkeypatch, [ALPHA, BETA], answer) + checked: list[str] = [] - class StubProject: - def __init__(self, name: str | None, external_id: str | None): - self.name = name or "main" - self.external_id = external_id or "local-main" + async def readiness(client, project): + checked.append(project.name) + return "# Project Index Required\n\nProject 'beta'" if project.name == "beta" else None - @asynccontextmanager - async def fake_get_project_client(project=None, context=None, project_id=None): - searched_projects.append((project, project_id)) - yield object(), StubProject(project, project_id) + monkeypatch.setattr(search_mod, "project_index_required", readiness) - async def fake_resolve_project_and_path(client, identifier, project=None, context=None): - return StubProject(project, None), identifier, False + result = await search_mod.search_notes( + query="missing", search_all_projects=True, output_format="text" + ) - class MockSearchClient: - def __init__(self, client, project_id): - self.project_id = project_id + assert isinstance(result, str) + assert result.startswith("# Project Index Required") + assert checked == ["alpha", "beta"] - async def search(self, payload, page, page_size): - return SearchResponse( - results=[ - SearchResult( - title=f"Note in {self.project_id or 'local'}", - permalink="notes/example", - content="", - type=SearchItemType.OBSERVATION - if compact_observations - else SearchItemType.ENTITY, - score=0.5, - file_path="/notes/example.md", - ) - ], - current_page=page, - page_size=page_size, - total=1, - ) - monkeypatch.setattr(search_mod, "_load_search_project_refs", fake_load_search_project_refs) - monkeypatch.setattr(search_mod, "get_project_client", fake_get_project_client) - monkeypatch.setattr(search_mod, "resolve_project_and_path", fake_resolve_project_and_path) - monkeypatch.setattr(clients_mod, "SearchClient", MockSearchClient) +@pytest.mark.asyncio +async def test_projects_selects_a_subset_across_databases(monkeypatch): + """Naming projects searches only those, grouped by the database each lives in.""" + search_mod = importlib.import_module("basic_memory.mcp.tools.search") + + def answer(workspace, project_ids, page, page_size): + ref = TEAM if workspace == TEAM.workspace_tenant_id else BETA + return _page([_result(ref, title=ref.name, score=0.5)], page, page_size) + + _workspaces, calls = _install_scoped_search(monkeypatch, [PERSONAL, TEAM, ALPHA, BETA], answer) result = await search_mod.search_notes( - query="anything", - search_all_projects=True, - output_format="json", - compact=compact_observations, + query="notes", projects=["beta", TEAM.external_id], output_format="json" ) assert isinstance(result, dict) - assert searched_projects == [("alpha", None), ("beta", None)], ( - "Local fan-out must omit project_id so the recursive search_notes calls " - "take the name-routed path." + assert [(call["workspace"], call["project_ids"]) for call in calls] == [ + (None, [BETA.id]), + (TEAM.workspace_tenant_id, [TEAM.id]), + ] + assert {item["title"] for item in result["results"]} == {"beta", "team-paul/main"} + + +@pytest.mark.asyncio +async def test_an_unknown_project_name_is_refused_before_any_query(monkeypatch): + search_mod = importlib.import_module("basic_memory.mcp.tools.search") + _workspaces, calls = _install_scoped_search( + monkeypatch, [ALPHA, BETA], lambda *args: pytest.fail("must not query") ) - assert result["total"] == 2 - assert result["total_is_exact"] is True - if compact_observations: - assert [item["permalink"] for item in result["results"]] == [ - "alpha/notes/example.md", - "beta/notes/example.md", - ] + + with pytest.raises(ValueError, match="Unknown project\\(s\\): gamma"): + await search_mod.search_notes(query="notes", projects=["gamma"], output_format="json") + + assert calls == [] + + +@pytest.mark.asyncio +async def test_a_hit_the_server_does_not_attribute_is_a_version_skew_error(monkeypatch): + """A result without project identity cannot be qualified; say so instead of guessing.""" + search_mod = importlib.import_module("basic_memory.mcp.tools.search") + + def answer(workspace, project_ids, page, page_size): + stray = _result(ALPHA, title="stray", score=0.5).model_copy( + update={"project_external_id": None} + ) + return _page([stray], page, page_size) + + _install_scoped_search(monkeypatch, [ALPHA, BETA], answer) + + # Like the valid-time skew above, this is a version mismatch, not a search miss: + # it is raised as an error rather than returned as an empty or partial page. + with pytest.raises(ValueError, match="did not attribute"): + await search_mod.search_notes(query="notes", search_all_projects=True, output_format="json") + + +def test_project_refs_need_an_id_an_external_id_and_a_name(): + """List rows a scoped search cannot address or attribute are left out.""" + search_mod = importlib.import_module("basic_memory.mcp.tools.search") + payload = { + "projects": [ + {"name": "alpha", "external_id": ALPHA.external_id, "id": 3, "path": "/alpha"}, + {"name": "no-id", "external_id": BETA.external_id}, + {"name": "no-external-id", "id": 5}, + {"name": "bool-id", "external_id": PERSONAL.external_id, "id": True}, + { + "name": "main", + "qualified_name": "team-paul/main", + "external_id": TEAM.external_id, + "id": 22, + "workspace_tenant_id": "tenant-team", + "path": "/team/main", + }, + ], + "constrained_project": None, + } + + refs = search_mod._search_project_refs(payload) + + assert refs == [ + SearchProjectRef( + name="alpha", + external_id=ALPHA.external_id, + id=3, + workspace_tenant_id=None, + path="/alpha", + ), + TEAM, + ] + assert search_mod._search_project_refs({"projects": "nope"}) == [] + assert search_mod._search_project_refs(None) == [] diff --git a/tests/mcp/tools/test_search_notes_multi_project_temporal.py b/tests/mcp/tools/test_search_notes_multi_project_temporal.py index 85fe04963..b0431d2b6 100644 --- a/tests/mcp/tools/test_search_notes_multi_project_temporal.py +++ b/tests/mcp/tools/test_search_notes_multi_project_temporal.py @@ -1,7 +1,7 @@ -"""All-projects search must carry the valid-time filter into every project (SPEC-82). +"""All-projects search must carry the valid-time filter into every database (SPEC-82). `_search_all_projects` re-declares the whole filter surface in its own signature and then -calls `search_notes` once per project. A filter that is not repeated there is dropped for +runs one scoped query per database. A filter that is not repeated there is dropped for every project at once, and the merged answer would quietly mix filtered and unfiltered rows -- the worst shape this failure can take, because the result still looks like an answer. @@ -10,52 +10,49 @@ import importlib from contextlib import asynccontextmanager from typing import Any +from unittest.mock import AsyncMock import pytest +from basic_memory.mcp.tools.search import SearchProjectRef from basic_memory.schemas.search import SearchItemType, SearchResponse, SearchResult PROJECT_REFS = [ - {"project": "personal/main", "project_id": "11111111-1111-1111-1111-111111111111"}, - {"project": "team-paul/main", "project_id": "22222222-2222-2222-2222-222222222222"}, + SearchProjectRef( + name="personal/main", + external_id="11111111-1111-1111-1111-111111111111", + id=11, + workspace_tenant_id="tenant-personal", + path="/personal/main", + ), + SearchProjectRef( + name="team-paul/main", + external_id="22222222-2222-2222-2222-222222222222", + id=22, + workspace_tenant_id="tenant-team", + path="/team/main", + ), ] -@pytest.fixture -def cloud_routing(monkeypatch): - """Pin the routing signals so project ids are forwarded deterministically.""" - search_mod = importlib.import_module("basic_memory.mcp.tools.search") - monkeypatch.setattr(search_mod, "is_factory_mode", lambda: False) - monkeypatch.setattr(search_mod, "_explicit_routing", lambda: True) - monkeypatch.setattr(search_mod, "_force_local_mode", lambda: False) - monkeypatch.setattr(search_mod, "has_cloud_credentials", lambda config: True) - - def _install_stub_client(monkeypatch, payloads: list[dict[str, Any]], refs) -> None: - """Route every per-project search into a stub that records its query payload.""" + """Route every database's scoped search into a stub that records its query payload.""" clients_mod = importlib.import_module("basic_memory.mcp.clients") search_mod = importlib.import_module("basic_memory.mcp.tools.search") - - class StubProject: - def __init__(self, name: str | None, external_id: str | None): - self.name = name or "main" - self.external_id = external_id or "local-main" + refs_by_tenant = {ref.workspace_tenant_id: ref for ref in refs} @asynccontextmanager - async def fake_get_project_client(project=None, context=None, project_id=None): - yield object(), StubProject(project, project_id) - - async def fake_resolve_project_and_path(client, identifier, project=None, context=None): - return StubProject(project, None), identifier, False + async def fake_get_client(project_name=None, workspace=None): + yield {"workspace": workspace} async def fake_load_search_project_refs(context=None): return refs - class MockSearchClient: - def __init__(self, client, project_id): - self.project_id = project_id + class StubScopedSearchClient: + def __init__(self, client): + self.ref = refs_by_tenant[client["workspace"]] - async def search(self, payload, page, page_size): + async def search(self, payload, *, project_ids, page, page_size): payloads.append(payload) return SearchResponse( results=[ @@ -66,6 +63,8 @@ async def search(self, payload, page, page_size): type=SearchItemType.OBSERVATION, score=0.5, file_path="/main/decisions/cache-layer.md", + project_id=self.ref.id, + project_external_id=self.ref.external_id, ) ], current_page=page, @@ -75,14 +74,14 @@ async def search(self, payload, page, page_size): ) monkeypatch.setattr(search_mod, "_load_search_project_refs", fake_load_search_project_refs) - monkeypatch.setattr(search_mod, "get_project_client", fake_get_project_client) - monkeypatch.setattr(search_mod, "resolve_project_and_path", fake_resolve_project_and_path) - monkeypatch.setattr(clients_mod, "SearchClient", MockSearchClient) + monkeypatch.setattr(search_mod, "get_client", fake_get_client) + monkeypatch.setattr(clients_mod, "ScopedSearchClient", StubScopedSearchClient) + monkeypatch.setattr(search_mod, "project_index_required", AsyncMock(return_value=None)) @pytest.mark.asyncio -async def test_all_projects_search_forwards_the_valid_time_filter(monkeypatch, cloud_routing): - """Every project is asked the same valid-time question, not just the first.""" +async def test_all_projects_search_forwards_the_valid_time_filter(monkeypatch): + """Every database is asked the same valid-time question, not just the first.""" search_mod = importlib.import_module("basic_memory.mcp.tools.search") payloads: list[dict[str, Any]] = [] _install_stub_client(monkeypatch, payloads, PROJECT_REFS) @@ -101,12 +100,12 @@ async def test_all_projects_search_forwards_the_valid_time_filter(monkeypatch, c assert payload["valid_at"] == "2026-07-28" assert payload["time_kind"] == "effective" assert payload["valid_overlaps"] is None - # Every leg confirmed it ran the filter, so the merged answer confirms it too. + # Every database confirmed it ran the filter, so the merged answer confirms it too. assert result["temporal_applied"] is True @pytest.mark.asyncio -async def test_all_projects_search_forwards_an_overlap_filter(monkeypatch, cloud_routing): +async def test_all_projects_search_forwards_an_overlap_filter(monkeypatch): payloads: list[dict[str, Any]] = [] _install_stub_client(monkeypatch, payloads, PROJECT_REFS) search_mod = importlib.import_module("basic_memory.mcp.tools.search") @@ -125,7 +124,7 @@ async def test_all_projects_search_forwards_an_overlap_filter(monkeypatch, cloud @pytest.mark.asyncio -async def test_all_projects_search_without_a_filter_claims_nothing(monkeypatch, cloud_routing): +async def test_all_projects_search_without_a_filter_claims_nothing(monkeypatch): """An ordinary all-projects search stays exactly the payload it always was.""" payloads: list[dict[str, Any]] = [] _install_stub_client(monkeypatch, payloads, PROJECT_REFS) @@ -150,9 +149,8 @@ async def test_all_projects_search_without_a_filter_claims_nothing(monkeypatch, ], ) @pytest.mark.asyncio -async def test_a_malformed_filter_is_refused_before_any_project_is_searched( +async def test_a_malformed_filter_is_refused_before_any_database_is_searched( monkeypatch, - cloud_routing, valid_at: str | None, valid_overlaps: str | None, time_kind: str | None, @@ -160,12 +158,10 @@ async def test_a_malformed_filter_is_refused_before_any_project_is_searched( ): """A typo must read as an error, never as an all-projects search with no matches. - Each per-project leg turns the API's 400 into a `# Search Failed` string, which the - fan-out cannot tell from a project being unavailable: it logs it and skips on. With - every project skipped the merged answer is an empty success that still reports - `temporal_applied`, so a mistyped filter would come back as the plausible-looking - "no matches found" for a question that never ran anywhere. Client-side validation is - the only layer that can tell the two apart, so it runs once, before the fan-out. + A database that refuses the filter is logged and skipped, so with every database + skipped the merged answer would be an empty success that still reports + `temporal_applied`. Client-side validation is the only layer that can tell a typo + from an unavailable database, so it runs once, before any query. """ payloads: list[dict[str, Any]] = [] _install_stub_client(monkeypatch, payloads, PROJECT_REFS) @@ -185,9 +181,7 @@ async def test_a_malformed_filter_is_refused_before_any_project_is_searched( @pytest.mark.asyncio -async def test_all_projects_search_with_no_projects_still_confirms_the_filter( - monkeypatch, cloud_routing -): +async def test_all_projects_search_with_no_projects_still_confirms_the_filter(monkeypatch): """Zero projects is an empty answer to the valid-time question, not an unfiltered one.""" payloads: list[dict[str, Any]] = [] _install_stub_client(monkeypatch, payloads, []) From ae5520a68c522fdbbbbaa6b98c0ec8df5987c001 Mon Sep 17 00:00:00 2001 From: phernandez Date: Tue, 15 Sep 2026 18:01:20 -0500 Subject: [PATCH 7/8] docs(mcp): regenerate the search-notes man page for the projects parameter `tests/test_man_pages.py` compares the rendered SYNOPSIS and PARAMETERS of search-notes(3) with the tool registry; the new `projects` parameter left the committed page stale. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_019YW9ysxugGGBCNEGzsxtFV Signed-off-by: phernandez --- src/basic_memory/man/man3/search-notes(3).md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/basic_memory/man/man3/search-notes(3).md b/src/basic_memory/man/man3/search-notes(3).md index 83804f61e..f252de77d 100644 --- a/src/basic_memory/man/man3/search-notes(3).md +++ b/src/basic_memory/man/man3/search-notes(3).md @@ -21,7 +21,7 @@ MCP: ``` search_notes(query=None, project=None, project_id=None, - search_all_projects=False, page=1, page_size=10, + search_all_projects=False, projects=None, page=1, page_size=10, search_type=None, output_format="text", note_types=None, entity_types=None, categories=None, after_date=None, metadata_filters=None, tags=None, status=None, @@ -88,6 +88,7 @@ It is not a hard token budget: titles and metadata can still be large. - **project** (string | null, optional, default: None) — Project name to search in. Optional - server will resolve using hierarchy. If unknown, use list_memory_projects() to discover available projects. - **project_id** (string | null, optional, default: None) — Project external_id (UUID). Prefer this over `project` when known — it routes to the exact project regardless of name collisions across cloud workspaces. Takes precedence over `project`. Get from list_memory_projects(). - **search_all_projects** (boolean, optional, default: False) — Optional opt-in to search every accessible project. Ignored when `project` or `project_id` is supplied. +- **projects** (array | null, optional, default: None) — Optional list of project names or external ids to search together. Names are matched exactly as list_memory_projects() reports them (cloud projects by their workspace-qualified name). Ignored when `project` or `project_id` is supplied; an unknown name is an error rather than a silent skip. - **page** (integer, optional, default: 1) — The page number of results to return (default 1). Aliases: page_number. - **page_size** (integer, optional, default: 10) — The number of results to return per page (default 10). Aliases: limit, per_page. - **search_type** (string | null, optional, default: None) — Type of search to perform, one of: "text", "title", "permalink", "vector", "semantic", "hybrid". Default is dynamic: "hybrid" when semantic search is enabled, otherwise "text". From e9c4bff93ce0f52f5dd84d4040b2a2faf5f1e0ce Mon Sep 17 00:00:00 2001 From: phernandez Date: Tue, 15 Sep 2026 17:53:56 -0500 Subject: [PATCH 8/8] perf(core): read only the projects in scope from the vector indexes Both built-in vector adapters took the k nearest vectors across the whole database and applied the scope afterwards. In a shared database a small project could get an under-filled or empty page for a query its own notes answered, because a larger neighbour's vectors filled the window first. On a 20k-chunk corpus where one project holds 2% of the vectors, a limit-100 query for that project returned 1 row from sqlite-vec. sqlite-vec: `project_id` becomes a vec0 partition key, so the KNN ranks each partition in scope and the outer ORDER BY merges them (100 of 100 rows, 0.2 ms against 3.1 ms). Existing local storage is carried into the partitioned table through a temporary table inside one transaction; vectors and ready manifests survive, nothing is re-embedded. Dimension and source_hash mismatches still rebuild as before. pgvector: the statement's tie-break sort keys had kept the planner off the HNSW index entirely, so every vector query was an exact scan of the table. The window is now taken by distance alone inside a materialized CTE and re-sorted with the tie-breaks outside it. `hnsw.ef_search` is set per transaction to the window the query must fill (the default 40 capped every larger candidate pool), and `hnsw.iterative_scan = relaxed_order` lets the scan continue until the scope and manifest filters have admitted enough rows. That scan needs pgvector 0.8; an older extension now raises SemanticDependenciesMissingError at first use rather than returning short windows. Neon tenants run 0.8.0. Refs #1558 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_019YW9ysxugGGBCNEGzsxtFV Signed-off-by: phernandez --- CHANGELOG.md | 13 ++ src/basic_memory/models/search.py | 9 +- src/basic_memory/repository/pgvector_index.py | 64 +++++++- .../repository/sqlite_vec_index.py | 60 +++++++- tests/repository/test_pgvector_index.py | 86 ++++++++++- .../test_postgres_search_repository_unit.py | 3 + .../test_sqlite_vector_search_repository.py | 142 +++++++++++++++++- 7 files changed, 361 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cc000c6df..6de2d3fc1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,19 @@ database is skipped with a warning and an inexact total, and a retryable outage or a server too old to attribute its results fails the whole page. +- **#1558**: Vector retrieval reads only the projects in scope and fills the window it + asks for. The sqlite-vec table gains a `project_id` partition key, so a scoped + nearest-neighbour query ranks each project's own vectors instead of taking the k + nearest across the whole database and discarding the out-of-scope ones, which + could leave a small project with an empty page for a query its notes answered. + Existing local storage is carried into the partitioned table without re-embedding. + On Postgres the nearest-neighbour statement now runs on the HNSW index (its + tie-break sort keys had kept the planner on an exact scan of every vector), with + `hnsw.ef_search` sized to the candidate window and an iterative scan that keeps + going until the scope and manifest filters have admitted enough rows. That scan + needs pgvector 0.8 or later; an older extension is reported as a dependency error + instead of quietly returning short windows. + - **#1512**: Word, PowerPoint, and CSV files get the same sidecar Markdown note a PDF gets. `bm import document ` indexes the project, extracts the file, and writes `..md` next to it plus a run note under diff --git a/src/basic_memory/models/search.py b/src/basic_memory/models/search.py index 70badd098..d192c5250 100644 --- a/src/basic_memory/models/search.py +++ b/src/basic_memory/models/search.py @@ -223,11 +223,18 @@ def create_sqlite_search_vector_embeddings(dimensions: int) -> DDL: - """Build sqlite-vec virtual table DDL for the configured embedding dimension.""" + """Build sqlite-vec virtual table DDL for the configured embedding dimension. + + ``project_id`` is a vec0 partition key: a scoped nearest-neighbour query reads + only the partitions in scope, instead of ranking every project's vectors in the + database and discarding the out-of-scope ones afterwards, which left a small + project with an under-filled window whenever a larger neighbour sat closer. + """ return DDL( f""" CREATE VIRTUAL TABLE IF NOT EXISTS search_vector_embeddings USING vec0( + project_id integer partition key, embedding float[{dimensions}], +source_hash text ) diff --git a/src/basic_memory/repository/pgvector_index.py b/src/basic_memory/repository/pgvector_index.py index 1db0be8c5..41ad25e31 100644 --- a/src/basic_memory/repository/pgvector_index.py +++ b/src/basic_memory/repository/pgvector_index.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import re from collections.abc import Sequence from loguru import logger @@ -23,6 +24,32 @@ ) +# pgvector's HNSW scan hands back at most ``hnsw.ef_search`` rows, 40 by default, +# whatever the LIMIT asks for; the server caps the setting at 1000. +HNSW_EF_SEARCH_DEFAULT = 40 +HNSW_EF_SEARCH_MAX = 1000 + +_PGVECTOR_VERSION = re.compile(r"^(\d+)\.(\d+)") + + +def pgvector_supports_iterative_scan(extversion: str) -> bool: + """Whether this pgvector keeps scanning an HNSW index until the LIMIT is filled. + + Iterative index scans arrived in pgvector 0.8.0. Before that, a scan stops at + ``hnsw.ef_search`` candidates, and a filter applied afterwards (the manifest + join, a scope narrower than the table) leaves the window under-filled, so the + adapter cannot promise the nearest ``limit`` rows in scope. A version string + the pattern cannot read counts as older. + """ + match = _PGVECTOR_VERSION.match(extversion) + return match is not None and (int(match.group(1)), int(match.group(2))) >= (0, 8) + + +def hnsw_ef_search_for(limit: int) -> int: + """The candidate-list size that lets one HNSW scan return ``limit`` rows.""" + return min(max(limit, HNSW_EF_SEARCH_DEFAULT), HNSW_EF_SEARCH_MAX) + + class PgVectorIndex: """Persist and query semantic vectors in PostgreSQL with pgvector.""" @@ -56,6 +83,19 @@ async def initialize(self) -> None: raise SemanticDependenciesMissingError( "pgvector extension is unavailable for this Postgres database." ) from exc + version = await session.execute( + text("SELECT extversion FROM pg_extension WHERE extname = 'vector'") + ) + extversion = str(version.scalar_one()) + # Trigger: the installed pgvector predates iterative index scans. + # Why: without them a scoped query silently returns fewer rows than + # the window asked for; a deployment gap should read as one. + # Outcome: a typed dependency error the API reports as a bad request. + if not pgvector_supports_iterative_scan(extversion): + raise SemanticDependenciesMissingError( + f"pgvector {extversion} predates iterative index scans; semantic " + "search needs pgvector 0.8 or later (ALTER EXTENSION vector UPDATE)." + ) existing_dimensions = await self._existing_dimensions(session) storage_missing = existing_dimensions is None @@ -322,10 +362,25 @@ async def search( embeddings_in_scope = projects.predicate("e.project_id", params) chunks_in_scope = projects.predicate("c.project_id", params) async with db.scoped_session(self._session_maker) as session: + # Both settings are transaction-local, so they last exactly as long as + # this scoped session. The scan is sized to the window it must fill + # and continues past that until the scope and manifest filters have + # admitted enough rows, instead of stopping at the first ef_search + # candidates and returning whatever of them survived. + await session.execute( + text("SELECT set_config('hnsw.ef_search', :ef_search, true)"), + {"ef_search": str(hnsw_ef_search_for(limit))}, + ) + await session.execute( + text("SELECT set_config('hnsw.iterative_scan', 'relaxed_order', true)") + ) + # A relaxed iterative scan may hand rows back slightly out of distance + # order, so the window is taken by distance alone and sorted once more. result = await session.execute( text( + "WITH nearest AS MATERIALIZED (" "SELECT c.entity_id, c.chunk_key, " - "1 - (e.embedding <=> CAST(:query AS vector)) AS similarity " + "e.embedding <=> CAST(:query AS vector) AS distance " "FROM search_vector_embeddings e " "JOIN search_vector_chunks c ON c.id = e.chunk_id " f"WHERE {embeddings_in_scope} " @@ -335,9 +390,12 @@ async def search( "AND c.embedding_status = 'ready' " "AND c.embedding_model = :embedding_identity " "AND e.source_hash = c.source_hash " - "ORDER BY e.embedding <=> CAST(:query AS vector), " - "c.entity_id ASC, c.chunk_key ASC " + "ORDER BY e.embedding <=> CAST(:query AS vector) " "LIMIT :limit" + ") " + "SELECT entity_id, chunk_key, 1 - distance AS similarity " + "FROM nearest " + "ORDER BY distance ASC, entity_id ASC, chunk_key ASC" ), params, ) diff --git a/src/basic_memory/repository/sqlite_vec_index.py b/src/basic_memory/repository/sqlite_vec_index.py index cfa45ddf7..9f81baf06 100644 --- a/src/basic_memory/repository/sqlite_vec_index.py +++ b/src/basic_memory/repository/sqlite_vec_index.py @@ -115,6 +115,7 @@ async def initialize(self) -> None: expected_dimensions = f"float[{self.scope.dimensions}]" dimensions_changed = bool(vector_sql and expected_dimensions not in vector_sql) source_hash_missing = bool(vector_sql and "+source_hash text" not in vector_sql) + partitions_missing = bool(vector_sql and "partition key" not in vector_sql) if dimensions_changed or source_hash_missing: logger.warning( "SQLite vector storage schema mismatch " @@ -124,6 +125,19 @@ async def initialize(self) -> None: source_hash_missing=source_hash_missing, ) await session.execute(text("DROP TABLE IF EXISTS search_vector_embeddings")) + elif partitions_missing: + # Trigger: storage predates the project_id partition key, and its + # vectors are otherwise current. + # Why: vec0 cannot add a column in place, and re-embedding a whole + # vault only to change how rows are partitioned would cost every + # local user a full embedding pass for nothing new. + # Outcome: the vectors are carried into partitioned storage, each + # keyed by the project its manifest row names; manifests stay ready. + logger.info( + "SQLite vector storage predates project partitions; " + "carrying vectors into partitioned storage" + ) + await self._partition_existing_storage(session) await session.execute(create_sqlite_search_vector_embeddings(self.scope.dimensions)) # Missing or dimension-rebuilt vec storage has no vectors, so ready @@ -138,6 +152,35 @@ async def initialize(self) -> None: await session.commit() self._initialized = True + async def _partition_existing_storage(self, session: AsyncSession) -> None: + """Rebuild vec storage with the partition key, keeping every current vector. + + SQLite DDL is transactional, so the copy out, drop, recreate, and copy back + either all land or none do. A vector whose manifest row is gone has no + project to file under and is left behind, which is what the orphan sweep + would have done to it anyway. + """ + await session.execute( + text( + "CREATE TEMP TABLE search_vector_embeddings_carry AS " + "SELECT e.rowid AS id, c.project_id AS project_id, " + "e.embedding AS embedding, e.source_hash AS source_hash " + "FROM search_vector_embeddings e " + "JOIN search_vector_chunks c ON c.id = e.rowid" + ) + ) + await session.execute(text("DROP TABLE search_vector_embeddings")) + await session.execute(create_sqlite_search_vector_embeddings(self.scope.dimensions)) + await session.execute( + text( + "INSERT INTO search_vector_embeddings " + "(rowid, project_id, embedding, source_hash) " + "SELECT id, project_id, embedding, source_hash " + "FROM search_vector_embeddings_carry" + ) + ) + await session.execute(text("DROP TABLE search_vector_embeddings_carry")) + async def upsert(self, project_id: int, records: Sequence[VectorRecord]) -> None: if not records: return @@ -192,12 +235,14 @@ async def upsert(self, project_id: int, records: Sequence[VectorRecord]) -> None ) await session.execute( text( - "INSERT INTO search_vector_embeddings (rowid, embedding, source_hash) " - "VALUES (:rowid, :embedding, :source_hash)" + "INSERT INTO search_vector_embeddings " + "(rowid, project_id, embedding, source_hash) " + "VALUES (:rowid, :project_id, :embedding, :source_hash)" ), [ { "rowid": rowids_by_key[record.key], + "project_id": project_id, "embedding": json.dumps(record.values), "source_hash": record.source_hash, } @@ -323,21 +368,24 @@ async def search( "embedding_identity": self.scope.embedding_identity, "limit": limit, } - chunks_in_scope = projects.predicate("c.project_id", params) + # vec0 ranks the k nearest within each partition in scope, so a small + # project is never crowded out of its own window by a larger neighbour that + # shares the database; the outer ORDER BY merges the partitions. + partitions_in_scope = projects.predicate("project_id", params) async with db.scoped_session(self._session_maker) as session: await self._ensure_loaded(session) result = await session.execute( text( "WITH vector_matches AS MATERIALIZED (" " SELECT rowid, distance, source_hash FROM search_vector_embeddings " - " WHERE embedding MATCH :query AND k = :vector_k" + f" WHERE {partitions_in_scope} " + " AND embedding MATCH :query AND k = :vector_k" ") " "SELECT c.entity_id, c.chunk_key, vector_matches.distance " "FROM vector_matches " "JOIN search_vector_chunks c ON c.id = vector_matches.rowid " "AND c.source_hash = vector_matches.source_hash " - f"WHERE {chunks_in_scope} " - "AND c.vector_index = 'sqlite-vec' " + "WHERE c.vector_index = 'sqlite-vec' " "AND c.embedding_status = 'ready' " "AND c.embedding_model = :embedding_identity " "ORDER BY vector_matches.distance ASC, " diff --git a/tests/repository/test_pgvector_index.py b/tests/repository/test_pgvector_index.py index c46abca42..ea9b5908e 100644 --- a/tests/repository/test_pgvector_index.py +++ b/tests/repository/test_pgvector_index.py @@ -9,7 +9,11 @@ import pytest from basic_memory.repository import pgvector_index as pgvector_index_module -from basic_memory.repository.pgvector_index import PgVectorIndex +from basic_memory.repository.pgvector_index import ( + PgVectorIndex, + hnsw_ef_search_for, + pgvector_supports_iterative_scan, +) from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError from basic_memory.repository.semantic_vector_index import ( @@ -40,6 +44,10 @@ def fetchone(self) -> object | None: def scalar_one_or_none(self) -> object | None: return self._scalar + def scalar_one(self) -> object: + assert self._scalar is not None + return self._scalar + def mappings(self) -> FakeResult: return self @@ -59,6 +67,7 @@ def __init__( chunk_rows: list[dict[str, object]] | None = None, search_rows: list[dict[str, object]] | None = None, fail_extension: bool = False, + pgvector_version: str = "0.8.0", ) -> None: self.table_exists = table_exists self.dimensions = dimensions @@ -66,6 +75,7 @@ def __init__( self.chunk_rows = chunk_rows or [] self.search_rows = search_rows or [] self.fail_extension = fail_extension + self.pgvector_version = pgvector_version self.calls: list[tuple[str, dict[str, object] | None]] = [] self.commit_count = 0 @@ -78,6 +88,8 @@ async def execute( self.calls.append((sql, params)) if "CREATE EXTENSION" in sql and self.fail_extension: raise RuntimeError("extension unavailable") + if "SELECT extversion" in sql: + return FakeResult(scalar=self.pgvector_version) if "information_schema.tables" in sql: return FakeResult(fetchone=(1,) if self.table_exists else None) if "attname = 'source_hash'" in sql: @@ -330,7 +342,7 @@ async def test_search_returns_normalized_stable_matches(monkeypatch) -> None: (15, 0.0), ] search_call = next(call for call in session.calls if "AS similarity" in call[0]) - assert "c.entity_id ASC, c.chunk_key ASC" in search_call[0] + assert "ORDER BY distance ASC, entity_id ASC, chunk_key ASC" in search_call[0] assert search_call[1] == { "query": "[1,0,0,0]", "scope_0": 7, @@ -356,3 +368,73 @@ async def test_search_binds_every_project_in_scope(monkeypatch) -> None: assert "AND c.project_id IN (:scope_0, :scope_1)" in sql assert params["scope_0"] == 7 assert params["scope_1"] == 9 + + +def _settings(session: FakeSession) -> list[tuple[str, dict[str, object] | None]]: + return [call for call in session.calls if "set_config" in call[0]] + + +@pytest.mark.parametrize( + ("extversion", "expected"), + [("0.7.4", False), ("0.8.0", True), ("0.8.1", True), ("1.0.0", True), ("garbage", False)], +) +def test_iterative_scan_arrived_in_pgvector_0_8(extversion: str, expected: bool) -> None: + assert pgvector_supports_iterative_scan(extversion) is expected + + +def test_ef_search_is_sized_to_the_window_within_the_server_bounds() -> None: + assert hnsw_ef_search_for(5) == 40 + assert hnsw_ef_search_for(40) == 40 + assert hnsw_ef_search_for(250) == 250 + assert hnsw_ef_search_for(5000) == 1000 + + +@pytest.mark.asyncio +async def test_initialize_requires_a_pgvector_that_can_keep_scanning(monkeypatch) -> None: + """An extension too old to fill a filtered window is a deployment error, not a quiet gap.""" + older = FakeSession(pgvector_version="0.7.4") + _install_session(monkeypatch, older) + index = PgVectorIndex(MagicMock(), _scope()) + + with pytest.raises(SemanticDependenciesMissingError, match="pgvector 0.7.4 predates"): + await index.initialize() + + assert not any("CREATE TABLE" in sql for sql in _sql_calls(older)) + assert index._initialized is False + + +@pytest.mark.asyncio +async def test_search_sizes_the_scan_to_the_window_it_must_fill(monkeypatch) -> None: + """A 250-row candidate window asks HNSW for 250 candidates, not the default 40.""" + session = FakeSession(search_rows=[]) + _install_session(monkeypatch, session) + index = PgVectorIndex(MagicMock(), _scope()) + index._initialized = True + + await index.search([1.0, 0.0, 0.0, 0.0], limit=250, projects=PROJECTS) + + ef_search = next(call for call in _settings(session) if "hnsw.ef_search" in call[0]) + assert "set_config('hnsw.ef_search', :ef_search, true)" in ef_search[0] + assert ef_search[1] == {"ef_search": "250"} + search_sql = next(sql for sql, _params in session.calls if "AS similarity" in sql) + # The window is taken by distance inside the CTE, then re-sorted with tie-breaks. + assert "ORDER BY e.embedding <=> CAST(:query AS vector) LIMIT :limit" in search_sql + + +@pytest.mark.asyncio +async def test_search_keeps_scanning_until_the_window_fills(monkeypatch) -> None: + session = FakeSession(search_rows=[]) + _install_session(monkeypatch, session) + index = PgVectorIndex(MagicMock(), _scope()) + index._initialized = True + + await index.search([1.0, 0.0, 0.0, 0.0], limit=5, projects=PROJECTS) + + settings = [sql for sql, _params in _settings(session)] + assert any("'hnsw.ef_search'" in sql for sql in settings) + assert any("'hnsw.iterative_scan', 'relaxed_order', true" in sql for sql in settings) + # Settings precede the scan they configure, inside the same session. + ordered = [sql for sql, _params in session.calls] + assert max(ordered.index(sql) for sql in settings) < next( + position for position, sql in enumerate(ordered) if "AS similarity" in sql + ) diff --git a/tests/repository/test_postgres_search_repository_unit.py b/tests/repository/test_postgres_search_repository_unit.py index 30a468e28..df18a17d8 100644 --- a/tests/repository/test_postgres_search_repository_unit.py +++ b/tests/repository/test_postgres_search_repository_unit.py @@ -225,10 +225,13 @@ async def fake_scoped_session(session_maker): ) missing_table = MagicMock() missing_table.fetchone.return_value = None + pgvector_version = MagicMock() + pgvector_version.scalar_one.return_value = "0.8.0" session.execute.side_effect = [ MagicMock(), MagicMock(), MagicMock(), + pgvector_version, missing_table, MagicMock(), MagicMock(), diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index c53291b26..9c0ec09b9 100644 --- a/tests/repository/test_sqlite_vector_search_repository.py +++ b/tests/repository/test_sqlite_vector_search_repository.py @@ -492,12 +492,12 @@ async def test_sqlite_vec_search_reads_every_project_in_scope(search_repository) ) await session.execute( text( - "INSERT INTO search_vector_embeddings (rowid, embedding, source_hash) " - "VALUES (:rowid, :embedding, 'hash')" + "INSERT INTO search_vector_embeddings (rowid, project_id, embedding, source_hash) " + "VALUES (:rowid, :project_id, :embedding, 'hash')" ), [ - {"rowid": 911, "embedding": "[1,0,0,0]"}, - {"rowid": 912, "embedding": "[0,1,0,0]"}, + {"rowid": 911, "project_id": own_project, "embedding": "[1,0,0,0]"}, + {"rowid": 912, "project_id": other_project, "embedding": "[0,1,0,0]"}, ], ) await session.commit() @@ -514,6 +514,140 @@ async def test_sqlite_vec_search_reads_every_project_in_scope(search_repository) assert nothing == [] +async def _seed_ready_vectors( + search_repository: SQLiteSearchRepository, + index: SQLiteVecIndex, + rows: list[tuple[int, int, str]], + *, + partitioned: bool = True, +) -> None: + """Insert ready manifest rows and their vectors: ``(rowid, project_id, embedding)``.""" + embedding_identity = search_repository._embedding_model_key() + async with db.scoped_session(search_repository.session_maker) as session: + await index._ensure_loaded(session) + await session.execute( + text( + "INSERT INTO search_vector_chunks (" + "id, entity_id, project_id, chunk_key, chunk_text, source_hash, " + "entity_fingerprint, embedding_model, vector_index, embedding_status" + ") VALUES (" + ":id, :id, :project_id, :chunk_key, 'text', 'hash', " + "'fingerprint', :embedding_model, 'sqlite-vec', 'ready')" + ), + [ + { + "id": rowid, + "project_id": project_id, + "chunk_key": f"entity:{rowid}:0", + "embedding_model": embedding_identity, + } + for rowid, project_id, _embedding in rows + ], + ) + if partitioned: + await session.execute( + text( + "INSERT INTO search_vector_embeddings " + "(rowid, project_id, embedding, source_hash) " + "VALUES (:rowid, :project_id, :embedding, 'hash')" + ), + [ + {"rowid": rowid, "project_id": project_id, "embedding": embedding} + for rowid, project_id, embedding in rows + ], + ) + else: + await session.execute( + text( + "INSERT INTO search_vector_embeddings (rowid, embedding, source_hash) " + "VALUES (:rowid, :embedding, 'hash')" + ), + [{"rowid": rowid, "embedding": embedding} for rowid, _project, embedding in rows], + ) + await session.commit() + + +@pytest.mark.asyncio +async def test_sqlite_vec_scope_is_a_partition_not_a_filter_on_the_nearest(search_repository): + """A small project fills its window even when a neighbour's vectors sit closer. + + The k nearest across the whole database used to be taken first and the scope + applied afterwards, so a project holding a few vectors among a large + neighbour's could get an empty page for a query its own notes answered. + """ + if not isinstance(search_repository, SQLiteSearchRepository): + pytest.skip("sqlite-vec search behavior is local SQLite-only.") + + _enable_semantic(search_repository) + await search_repository.init_search_index() + index = cast(SQLiteVecIndex, search_repository._semantic_vector_index) + small = search_repository.project_id + large = small + 1 + + # The large project's vectors are all nearer the query than the small one's. + await _seed_ready_vectors( + search_repository, + index, + [(921, large, "[1,0,0,0]"), (922, large, "[0.9,0.1,0,0]"), (923, large, "[0.8,0.2,0,0]")] + + [(931, small, "[0,1,0,0]"), (932, small, "[0,0,1,0]")], + ) + + nearest_two = await index.search( + [1.0, 0.0, 0.0, 0.0], limit=2, projects=ProjectScope.single(small) + ) + + assert [match.key.entity_id for match in nearest_two] == [931, 932] + + +@pytest.mark.asyncio +async def test_sqlite_vec_partitions_legacy_storage_without_re_embedding(search_repository): + """Storage from before the partition key is carried over, vectors and readiness intact.""" + if not isinstance(search_repository, SQLiteSearchRepository): + pytest.skip("sqlite-vec storage upgrade is local SQLite-only.") + + _enable_semantic(search_repository) + await search_repository.init_search_index() + index = cast(SQLiteVecIndex, search_repository._semantic_vector_index) + project = search_repository.project_id + dimensions = search_repository._vector_dimensions + + async with db.scoped_session(search_repository.session_maker) as session: + await index._ensure_loaded(session) + await session.execute(text("DROP TABLE search_vector_embeddings")) + await session.execute( + text( + "CREATE VIRTUAL TABLE search_vector_embeddings USING vec0(" + f"embedding float[{dimensions}], +source_hash text)" + ) + ) + await session.commit() + await _seed_ready_vectors( + search_repository, + index, + [(941, project, "[1,0,0,0]"), (942, project, "[0,1,0,0]")], + partitioned=False, + ) + + index.invalidate_initialization() + await index.initialize() + + async with db.scoped_session(search_repository.session_maker) as session: + table_sql = await session.scalar( + text("SELECT sql FROM sqlite_master WHERE name = 'search_vector_embeddings'") + ) + statuses = await session.execute( + text("SELECT embedding_status FROM search_vector_chunks WHERE id IN (941, 942)") + ) + carried = await session.execute( + text("SELECT rowid, project_id FROM search_vector_embeddings ORDER BY rowid") + ) + assert table_sql is not None and "project_id integer partition key" in table_sql + assert statuses.scalars().all() == ["ready", "ready"] + assert carried.all() == [(941, project), (942, project)] + found = await index.search([1.0, 0.0, 0.0, 0.0], limit=5, projects=ProjectScope.single(project)) + assert [match.key.entity_id for match in found] == [941, 942] + + @pytest.mark.asyncio async def test_sqlite_vec_delete_requires_pending_source_generation(search_repository): """A stale delete cannot remove a same-source vector that is already ready."""