diff --git a/CHANGELOG.md b/CHANGELOG.md index 239f084dc..6de2d3fc1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,38 @@ ### 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`. + +- **#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. + +- **#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 @@ -109,6 +141,49 @@ 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. + +- **#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. + +- **#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. + +- **#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/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/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". 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/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/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/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/pgvector_index.py b/src/basic_memory/repository/pgvector_index.py index 764a1c716..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 @@ -10,6 +11,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, @@ -22,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.""" @@ -55,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 @@ -151,37 +192,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 +200,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 +235,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 +263,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 +305,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 +314,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 +335,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,36 +346,58 @@ 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: + # 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 " - "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 " "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" ), - { - "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_query.py b/src/basic_memory/repository/postgres_search_query.py new file mode 100644 index 000000000..0fa2fc339 --- /dev/null +++ b/src/basic_memory/repository/postgres_search_query.py @@ -0,0 +1,861 @@ +"""PostgreSQL tsquery preparation and execution. + +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 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, + 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.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) +_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 = ("&", "|", "!", "(", ")", ":") + +# 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 --- + + +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, + query: PreparedSearchQuery, + *, + 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, 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 ("", "*"): + 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 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 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") + 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 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 + # 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 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 a6d7e3088..b57b4eb01 100644 --- a/src/basic_memory/repository/postgres_search_repository.py +++ b/src/basic_memory/repository/postgres_search_repository.py @@ -2,13 +2,9 @@ import asyncio import json -import re -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 @@ -21,28 +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, relaxation_word_tokens -from basic_memory.repository.script_ngrams import analyze_script_query, build_script_ngrams +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.temporal_filters import build_temporal_predicate +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 ( @@ -55,113 +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 - - -_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]: @@ -200,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 @@ -238,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 @@ -396,236 +267,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) # ------------------------------------------------------------------ @@ -669,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 @@ -727,22 +364,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.""" @@ -861,15 +482,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()" @@ -952,731 +564,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) - # ------------------------------------------------------------------ - - @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, - 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) --- - ( - from_clause, - where_clause, - params, - order_by_clause, - score_expr, - ) = await self._build_fts_query_parts( - 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, - ) - - # set limit and offset - 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, - {score_expr} as score - FROM {from_clause} - WHERE {where_clause} - ORDER BY score DESC {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 = self._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")): - 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 self._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, - ) - - ( - from_clause, - where_clause, - params, - _order_by_clause, - _score_expr, - ) = await self._build_fts_query_parts( - 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, - ) - sql = f"SELECT COUNT(*) FROM {from_clause} WHERE {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 = self._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")): - 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 self._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 new file mode 100644 index 000000000..2431d11cb --- /dev/null +++ b/src/basic_memory/repository/search_filters.py @@ -0,0 +1,309 @@ +"""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. ``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 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_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.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: + """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 + + +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, + query: PreparedSearchQuery, + candidate_keys: Sequence[SearchIndexKey] | 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 query.permalink: + params["permalink"] = query.permalink + conditions.append("search_index.permalink = :permalink") + + subtree_condition = file_path_prefix_condition(query.file_path_prefix, params) + if subtree_condition is not None: + conditions.append(subtree_condition) + + if candidate_keys is not None: + conditions.append(candidate_key_restriction_condition(candidate_keys, params)) + + if query.search_item_types: + type_placeholders: list[str] = [] + 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}") + 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 query.categories: + category_placeholders: list[str] = [] + for index, category in enumerate(query.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 query.note_types: + conditions.append( + build_note_type_predicate( + 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 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 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..59dd16182 100644 --- a/src/basic_memory/repository/search_query.py +++ b/src/basic_memory/repository/search_query.py @@ -2,6 +2,58 @@ 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 + + @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..8d5d80b28 --- /dev/null +++ b/src/basic_memory/repository/search_reader.py @@ -0,0 +1,1145 @@ +"""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, projects=self.scope + ) + 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, projects=self.scope + ) + 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. + """ + # 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) + 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. + """ + if self.scope.is_empty: + return [] + 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.py b/src/basic_memory/repository/search_repository.py index 2b9828202..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 @@ -235,7 +246,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, @@ -266,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/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 7e824a86e..8935b307d 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,28 +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_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, - 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, @@ -65,7 +59,6 @@ SemanticVectorIndexReconciler, VectorDeletion, VectorKey, - VectorMatch, VectorRecord, ) from basic_memory.repository.semantic_vector_sync import ( @@ -77,54 +70,19 @@ 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 -# --- 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. -# 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'" -) -# 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 +# --- Semantic sync constants --- + 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"}) -# 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"] @@ -174,158 +132,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, @@ -365,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 @@ -389,6 +196,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. @@ -405,6 +214,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,24 +266,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, search_text: Optional[str] = None, @@ -491,42 +284,42 @@ 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 - - Backend-specific implementations: - - SQLite: Uses MATCH operator and bm25() for scoring - - Postgres: Uses @@ operator and ts_rank() for scoring + 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``. """ - pass + 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, + ) + reader = await self._reader_for(query) + return await reader.search( + query, + limit=limit, + offset=offset, + allow_relaxed=allow_relaxed, + session=session, + candidate_keys=candidate_keys, + trace=trace, + ) async def count( self, @@ -545,10 +338,66 @@ async def count( min_similarity: Optional[float] = None, allow_relaxed: bool = False, ) -> int: - """Count results when a backend-specific COUNT query is available.""" - 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.") + """Count full-text matches with the same filters as ``search``.""" + 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 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) @@ -559,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.project_id, - 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.project_id, - 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, *, @@ -714,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, object] = { - "project_id": self.project_id, - "vector_index": self._semantic_vector_index_name, - "embedding_model": self._embedding_model_key(), - } - 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 " + CURRENT_VECTOR_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.project_id, 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, @@ -867,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: @@ -969,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, @@ -980,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, @@ -1146,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() @@ -1182,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 " @@ -1194,8 +846,8 @@ async def _finalize_prepared_vector_deletions( await session.commit() return - await self._semantic_vector_index.delete(deletions) - if self._semantic_vector_index_name in _BUILT_IN_VECTOR_INDEX_NAMES: + 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: await session.execute( @@ -1208,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 # ------------------------------------------------------------------ @@ -1512,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 @@ -1567,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.""" @@ -1610,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") @@ -1684,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: @@ -1725,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: @@ -1837,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 @@ -2174,950 +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)}, - "project_id": self.project_id, - } - 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 - 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/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/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/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_query.py b/src/basic_memory/repository/sqlite_search_query.py new file mode 100644 index 000000000..42dfdbd67 --- /dev/null +++ b/src/basic_memory/repository/sqlite_search_query.py @@ -0,0 +1,617 @@ +"""SQLite FTS5 query preparation and execution. + +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 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, + 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.repository.search_trace import SearchTraceCollector, build_fts_page_stage + +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)" + +# 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 --- + + +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, + query: PreparedSearchQuery, + *, + entity_columns: Collection[str], + 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, 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 ("", "*"): + 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 query.title: + params["title_text"] = prepare_search_term(query.title.strip(), is_prefix=False) + match_conditions.append("search_index.title MATCH :title_text") + + if query.permalink_match: + # GLOB patterns keep their syntax; prepare_search_term would quote the slashes. + permalink_text = query.permalink_match.lower().strip() + params["permalink"] = permalink_text + if "*" in query.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 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 + # 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 query.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 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 4160d7932..7f50b8df5 100644 --- a/src/basic_memory/repository/sqlite_search_repository.py +++ b/src/basic_memory/repository/sqlite_search_repository.py @@ -1,14 +1,10 @@ """SQLite FTS5-based search repository implementation.""" import asyncio -import re -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 @@ -27,36 +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, - 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.temporal_filters import build_temporal_predicate +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 - - -SQLITE_WORD_COLUMNS = "{title content_stems content_snippet}" class SQLiteSearchRepository(SearchRepositoryBase): @@ -80,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 @@ -110,20 +85,9 @@ 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), ) - 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. @@ -160,281 +124,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. @@ -576,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 @@ -641,25 +326,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, @@ -733,15 +399,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): @@ -775,608 +432,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) - # ------------------------------------------------------------------ - - @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, - 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) --- - ( - from_clause, - where_clause, - params, - order_by_clause, - score_expression, - ) = await self._build_fts_query_parts( - 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, - ) - - # set limit on search query - 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, - {score_expression} as score - FROM {from_clause} - WHERE {where_clause} - ORDER BY score ASC {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 = ( - self._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 self._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, - ) - - ( - from_clause, - where_clause, - params, - _order_by_clause, - _score_expression, - ) = await self._build_fts_query_parts( - 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, - ) - sql = f"SELECT COUNT(*) FROM {from_clause} WHERE {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 = ( - self._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 self._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_vec_index.py b/src/basic_memory/repository/sqlite_vec_index.py index c3344c31c..9f81baf06 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, @@ -78,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()")) @@ -103,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 " @@ -112,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 @@ -126,43 +152,43 @@ 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( + 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( - "SELECT id, entity_id, chunk_key FROM search_vector_chunks " - "WHERE project_id = :project_id AND (" + " OR ".join(predicates) + ")" - ), - params, + "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" + ) ) - return { - VectorKey(entity_id=int(row["entity_id"]), chunk_key=str(row["chunk_key"])): int( - row["id"] + 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" ) - for row in result.mappings().all() - } + ) + await session.execute(text("DROP TABLE search_vector_embeddings_carry")) - 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): @@ -209,12 +235,14 @@ async def upsert(self, 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, } @@ -223,13 +251,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 +292,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 +302,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 +344,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,38 +355,43 @@ 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, + } + # 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 " - "WHERE c.project_id = :project_id " - "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, " "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/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/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/project_readiness.py b/src/basic_memory/services/project_readiness.py index f37a14fc2..006bf0899 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_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, ) @@ -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/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/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 301e643f2..1ba056858 100644 --- a/src/basic_memory/services/search_service.py +++ b/src/basic_memory/services/search_service.py @@ -4,13 +4,14 @@ 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 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 @@ -23,7 +24,8 @@ 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_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 @@ -42,25 +44,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. @@ -159,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. @@ -215,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: @@ -305,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, @@ -501,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/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/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/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/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/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, []) 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 9e49048f7..dd5e14dc4 100644 --- a/tests/repository/test_hybrid_fusion.py +++ b/tests/repository/test_hybrid_fusion.py @@ -6,24 +6,26 @@ 3. Produces zero fused score when the source score is zero """ -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 @@ -49,161 +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 - - @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 + ``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 - def _prepare_search_term(self, term, is_prefix=True): - return term # pragma: no cover + 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, - allow_relaxed: bool = False, + 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_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..ea9b5908e 100644 --- a/tests/repository/test_pgvector_index.py +++ b/tests/repository/test_pgvector_index.py @@ -9,7 +9,12 @@ 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 ( VectorDeletion, @@ -39,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 @@ -58,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 @@ -65,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 @@ -77,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: @@ -93,10 +106,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 +224,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 +263,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 +282,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 +295,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,23 +328,113 @@ 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), (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]", - "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 + + +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_quoted_queries.py b/tests/repository/test_postgres_search_quoted_queries.py index a42f709f5..3f59f562b 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_query as postgres_search_query_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_query_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) + # PostgresFts binds the classifier at import; patch it where it is read. + monkeypatch.setattr( + postgres_search_query_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..197571fd6 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 +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, @@ -261,35 +269,32 @@ 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 "&" 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, PreparedSearchQuery(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 "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, + PreparedSearchQuery(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, 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" now = datetime.now(timezone.utc) rows = [ @@ -374,7 +379,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 +429,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 +1202,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 +1276,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_query_module.is_tsquery_syntax_error def record_syntax_error(exc: Exception) -> bool: is_syntax_error = real_is_syntax_error(exc) @@ -1275,7 +1284,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) + # PostgresFts binds the classifier at import; patch it where it is read. + monkeypatch.setattr( + postgres_search_query_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 +636,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_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 885cd96ba..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 @@ -17,7 +18,9 @@ from basic_memory import db from basic_memory.models.knowledge import Entity from basic_memory.repository.search_index_row import SearchIndexRow -from basic_memory.repository.search_repository_base import file_path_prefix_condition +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() ] ), @@ -335,7 +344,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_reader.py b/tests/repository/test_search_reader.py new file mode 100644 index 000000000..de29903cb --- /dev/null +++ b/tests/repository/test_search_reader.py @@ -0,0 +1,300 @@ +"""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, projects=SCOPE) + + +@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] + + +@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_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_search_trace.py b/tests/repository/test_search_trace.py index aa6f95a89..055050960 100644 --- a/tests/repository/test_search_trace.py +++ b/tests/repository/test_search_trace.py @@ -1,5 +1,6 @@ """Execution-native search trace builders and repository integration.""" +from basic_memory.repository.search_scope import ProjectScope from collections.abc import Sequence from dataclasses import replace from datetime import datetime, timezone @@ -13,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, @@ -67,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, ) @@ -79,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] @@ -524,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 @@ -685,10 +687,10 @@ def record_statement(_conn, _cursor, statement, _parameters, _context, _executem assert all("auth retrieval" not in match.chunk_key for match in collector.vector.chunk_matches) async with db.scoped_session(repository.session_maker) as session: - assert await classify_hydration_drops(session, repository.project_id, ()) == () + assert await classify_hydration_drops(session, repository.scope, ()) == () readiness_race = await classify_hydration_drops( session, - repository.project_id, + repository.scope, ( HydrationDropKey( entity_id=1, @@ -722,7 +724,7 @@ async def test_classify_hydration_drops_batches_large_unhealthy_candidate_set( async with db.scoped_session(session_maker) as session: classified = await classify_hydration_drops( session, - test_project.id, + ProjectScope.single(test_project.id), dropped_keys, ) @@ -750,7 +752,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( @@ -764,7 +766,7 @@ async def test_classify_hydration_drop_observes_pending_to_ready_transition( async with db.scoped_session(session_maker) as session: classified = await classify_hydration_drops( session, - test_project.id, + ProjectScope.single(test_project.id), ( HydrationDropKey( entity_id=1, diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 6e08fe31c..43fe5806f 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -1,21 +1,28 @@ """Tests for semantic search orchestration in SearchRepositoryBase.""" +from sqlalchemy.ext.asyncio import AsyncSession +from basic_memory.repository.search_scope import ProjectScope import asyncio import hashlib from collections.abc import Sequence 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, ) @@ -35,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 --- @@ -64,6 +72,8 @@ def __init__(self): # Bypass parent __init__ since we don't need a real session_maker for unit tests self.session_maker = None self.project_id = 1 + self.scope = ProjectScope.single(1) + self._fts = FakeFts() @override async def init_search_index(self): @@ -80,10 +90,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, @@ -103,6 +109,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, @@ -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,17 +142,12 @@ 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.""" scope = VectorIndexScope( namespace="basic-memory-test", - project_id=1, embedding_identity="stub:4", dimensions=4, ) @@ -167,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( @@ -181,16 +172,30 @@ async def search( query: Sequence[float], *, limit: int, + projects: ProjectScope, ) -> list[VectorMatch]: 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] @@ -472,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() @@ -505,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")] ) @@ -651,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" @@ -694,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", @@ -909,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" @@ -945,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_semantic_vector_sync.py b/tests/repository/test_semantic_vector_sync.py index 307b53030..0e1003424 100644 --- a/tests/repository/test_semantic_vector_sync.py +++ b/tests/repository/test_semantic_vector_sync.py @@ -1,5 +1,7 @@ """Focused edge-case coverage for shared semantic vector synchronization.""" +from sqlalchemy.ext.asyncio import AsyncSession +from basic_memory.repository.search_scope import ProjectScope import hashlib from collections.abc import Sequence from contextlib import asynccontextmanager @@ -36,6 +38,7 @@ class _TestRepository(SearchRepositoryBase): def __init__(self): self.session_maker = None self.project_id = 1 + self.scope = ProjectScope.single(1) @override async def init_search_index(self): @@ -45,10 +48,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, @@ -68,6 +67,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, @@ -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_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) diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index f981fac1c..9c0ec09b9 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,205 @@ 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, project_id, embedding, source_hash) " + "VALUES (:rowid, :project_id, :embedding, 'hash')" + ), + [ + {"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() + + 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 == [] + + +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.""" @@ -486,7 +686,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 +699,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 +1659,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_filter_candidate_restriction.py b/tests/repository/test_vector_filter_candidate_restriction.py index bb6549f6f..8407a037a 100644 --- a/tests/repository/test_vector_filter_candidate_restriction.py +++ b/tests/repository/test_vector_filter_candidate_restriction.py @@ -29,11 +29,15 @@ from basic_memory import db from basic_memory.repository.embedding_provider import EmbeddingProvider -from basic_memory.repository.search_repository_base import ( +from basic_memory.repository.search_filters import candidate_key_restriction_condition +from basic_memory.repository.search_query import PreparedSearchQuery +from basic_memory.repository.search_reader import ( VECTOR_FILTER_SCAN_LIMIT, VECTOR_HYDRATION_BATCH_SIZE, - candidate_key_restriction_condition, + 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_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, ) diff --git a/tests/repository/test_vector_pagination.py b/tests/repository/test_vector_pagination.py index 9139f4859..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 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 - - @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 - def _prepare_search_term(self, term, is_prefix=True): - return term # pragma: no cover - - @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, - *, - 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 1c6323aac..4314176d5 100644 --- a/tests/repository/test_vector_threshold.py +++ b/tests/repository/test_vector_threshold.py @@ -1,25 +1,22 @@ """Tests for semantic_min_similarity threshold filtering in vector search.""" -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 @@ -33,176 +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 - - # --- 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 - def _prepare_search_term(self, term, is_prefix=True): - return term # pragma: no cover - - @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, - *, - 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]) + semantic = vector_semantic(min_similarity=0.6) - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) + # 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)) - 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) - - # 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 + semantic = vector_semantic(min_similarity=0.8) + fetch_rows = AsyncMock() - # All scores below 0.8 - fake_rows = _make_vector_rows([0.5, 0.4, 0.3]) - - 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 + semantic = vector_semantic() - 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) - - 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, )