diff --git a/src/basic_memory/api/v2/routers/inspect_router.py b/src/basic_memory/api/v2/routers/inspect_router.py index d2c711e2f..185c8452b 100644 --- a/src/basic_memory/api/v2/routers/inspect_router.py +++ b/src/basic_memory/api/v2/routers/inspect_router.py @@ -4,12 +4,22 @@ from fastapi import APIRouter, HTTPException +from basic_memory.api.v2.utils import get_entities_by_id_lookup from basic_memory.deps import ( + EntityServiceV2ExternalDep, FileServiceV2ExternalDep, LinkResolverV2ExternalDep, ProjectExternalIdPathDep, SearchRepositoryV2ExternalDep, + SearchServiceV2ExternalDep, ) +from basic_memory.repository.semantic_errors import ( + RerankProviderContractError, + RerankTransientError, + SemanticDependenciesMissingError, + SemanticSearchDisabledError, +) +from basic_memory.repository.search_trace import HybridQueryTrace, VectorQueryTrace from basic_memory.schemas.inspect import ( InspectChunk, InspectChunkReadiness, @@ -17,8 +27,11 @@ InspectChunksResponse, InspectDetachedSearchRow, InspectIndexBehindRowsDetail, + InspectQueryRequest, + InspectQueryResponse, InspectRowsBehindFileDetail, InspectSearchRow, + query_trace_response, ) from basic_memory.services.retrieval_inspect import ( ChunkFresh, @@ -26,12 +39,50 @@ ChunkNotIndexed, ChunkIndexBehindRows, ChunkRowsBehindFile, + explain_query, inspect_entity_chunks, ) router = APIRouter(prefix="/inspect", tags=["inspect"]) +@router.post("/query", response_model=InspectQueryResponse) +async def inspect_query( + data: InspectQueryRequest, + project_id: ProjectExternalIdPathDep, + entity_service: EntityServiceV2ExternalDep, + search_service: SearchServiceV2ExternalDep, +) -> InspectQueryResponse: + """Run one search and return the trace captured by that exact execution.""" + del project_id # Route resolution scopes the injected search service. + try: + trace = await explain_query( + search_service, + data.query, + limit=data.limit, + offset=data.offset, + ) + except (SemanticSearchDisabledError, SemanticDependenciesMissingError, ValueError) as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + except RerankTransientError as exc: + raise HTTPException(status_code=503, detail=str(exc)) from exc + except RerankProviderContractError as exc: + raise HTTPException(status_code=502, detail=str(exc)) from exc + + # Nullable-owner rows stay inspectable; only known owners get external-id enrichment. + entity_ids = {result.entity_id for result in trace.final if result.entity_id is not None} + if isinstance(trace, (VectorQueryTrace, HybridQueryTrace)): + entity_ids.update(rejection.entity_id for rejection in trace.vector.drops) + entity_ids.update( + match.entity_id for match in trace.vector.chunk_matches if match.entity_id is not None + ) + entities_by_id = await get_entities_by_id_lookup(entity_service, sorted(entity_ids)) + external_ids_by_entity_id = { + entity_id: str(entity.external_id) for entity_id, entity in entities_by_id.items() + } + return query_trace_response(trace, external_ids_by_entity_id) + + @router.post("/chunks", response_model=InspectChunksResponse) async def inspect_chunks( data: InspectChunksRequest, diff --git a/src/basic_memory/api/v2/utils.py b/src/basic_memory/api/v2/utils.py index 4f3dfd951..853f2c76f 100644 --- a/src/basic_memory/api/v2/utils.py +++ b/src/basic_memory/api/v2/utils.py @@ -32,6 +32,17 @@ class EntityServiceBatchLookup(Protocol): async def get_entities_by_id(self, ids: List[int]) -> Sequence[Any]: ... +async def get_entities_by_id_lookup( + entity_service: EntityServiceBatchLookup, + entity_ids: Sequence[int], +) -> dict[int, Any]: + """Fetch an entity batch once and index it by internal identity.""" + if not entity_ids: + return {} + entities = await entity_service.get_entities_by_id(list(entity_ids)) + return {entity.id: entity for entity in entities} + + def _required_str(value: str | None, field_name: str) -> str: """Return a required search field or fail before producing invalid response data.""" if value is None: @@ -233,9 +244,10 @@ async def to_search_results( phase="fetch_entities", result_count=len(all_entity_ids), ): - if all_entity_ids: - entities = await entity_service.get_entities_by_id(list(all_entity_ids)) - entities_by_id = {e.id: e for e in entities} + entities_by_id = await get_entities_by_id_lookup( + entity_service, + list(all_entity_ids), + ) search_results = [] with logfire.span( diff --git a/src/basic_memory/cli/commands/inspect.py b/src/basic_memory/cli/commands/inspect.py index 7638053c6..b8702ed7b 100644 --- a/src/basic_memory/cli/commands/inspect.py +++ b/src/basic_memory/cli/commands/inspect.py @@ -1,5 +1,7 @@ """Read-only retrieval inspection commands.""" +from enum import Enum +from itertools import groupby from typing import Annotated, Optional, assert_never import typer @@ -20,9 +22,12 @@ InspectChunksResponse, InspectDetachedSearchRow, InspectIndexBehindRowsDetail, + InspectQueryCandidate, + InspectQueryResponse, InspectRowsBehindFileDetail, InspectSearchRow, ) +from basic_memory.schemas.search import SearchQuery, SearchRetrievalMode inspect_app = typer.Typer() app.add_typer(inspect_app, name="inspect", help="Inspect retrieval projections") @@ -30,6 +35,12 @@ console = Console() +class InspectQueryMode(str, Enum): + TEXT = "text" + VECTOR = "vector" + HYBRID = "hybrid" + + async def run_inspect_chunks( identifier: str, *, @@ -46,6 +57,26 @@ async def run_inspect_chunks( ) +async def run_inspect_query( + query: SearchQuery, + *, + limit: int, + offset: int, + project: str | None, + project_id: str | None, +) -> InspectQueryResponse: + """Resolve the project route and execute one traced query.""" + async with get_project_client(project=project, project_id=project_id) as ( + http_client, + active_project, + ): + return await InspectClient(http_client, active_project.external_id).inspect_query( + query, + limit=limit, + offset=offset, + ) + + def _text_preview(text: str, limit: int = 120) -> str: """Return a compact single-line chunk preview for human output.""" single_line = " ".join(text.split()) @@ -265,6 +296,319 @@ def _plain_chunks(response: InspectChunksResponse) -> None: ) +def _score_text(candidate: InspectQueryCandidate) -> str: + score = candidate.scores.final_score + return f"{score:.4f}" if score is not None else "-" + + +def _movement_text(candidate: InspectQueryCandidate) -> str: + before = candidate.scores.pre_rerank_rank + after = candidate.scores.post_rerank_rank + if before is None or after is None: + return "-" + movement = before - after + return f"{movement:+d}" + + +def _query_candidate_label(candidate: InspectQueryCandidate) -> str: + if candidate.title or candidate.permalink: + return candidate.title or candidate.permalink or "" + if candidate.type is not None and candidate.id is not None: + return f"{candidate.type}:{candidate.id}" + if candidate.rejection_detail is not None and candidate.rejection_detail.chunk_key: + return candidate.rejection_detail.chunk_key + return "unknown candidate" + + +def _query_candidate_id(candidate: InspectQueryCandidate) -> str | None: + if candidate.external_id is not None: + return candidate.external_id + if candidate.type is not None and candidate.id is not None: + return f"{candidate.type}:{candidate.id}" + return None + + +def _display_query( + response: InspectQueryResponse, + *, + show_misses: bool, + show_ids: bool, +) -> None: + """Render the traced query as a compact human-readable query plan.""" + header = Text() + header.append(f"{response.query}\n", style="bold cyan") + header.append(f"Mode: {response.retrieval_mode.value}") + header.append(f" · Project: {response.project_id}") + header.append( + f" · Window: {response.window.offset + 1}-" + f"{response.window.offset + response.window.limit}" + ) + console.print(Panel(header, title="Retrieval query", expand=False)) + + engine = response.engine + reranker = engine.reranker + engine_text = Text() + engine_text.append(f"Vector index: {engine.vector_index}\n") + engine_text.append(f"Embedding model: {engine.embedding_model}\n") + if engine.ready_rows is None: + engine_text.append("Readiness: n/a (FTS execution reads no vector manifest)\n") + else: + engine_text.append( + "Readiness: " + f"{engine.ready_rows} ready, {engine.pending_rows} pending, " + f"{engine.other_identity_rows} other identity\n" + ) + engine_text.append(f"Fusion: {engine.fusion_formula}\n") + engine_text.append( + f"Minimum similarity: {engine.min_similarity:.4f} ({engine.min_similarity_source})\n" + ) + if reranker.enabled: + status = "applied" if reranker.applied else f"skipped: {reranker.skipped_reason}" + engine_text.append( + f"Reranker: {reranker.model} · {reranker.candidates} candidates · {status}" + ) + else: + engine_text.append("Reranker: disabled") + console.print(Panel(engine_text, title="Engine", expand=False)) + + stage_table = Table(title="Stages", show_header=True, header_style="bold") + stage_table.add_column("Stage") + stage_table.add_column("In", justify="right") + stage_table.add_column("Out", justify="right") + stage_table.add_column("Dropped", justify="right") + stage_table.add_column("ms", justify="right") + for stage in response.stages: + stage_name = stage.name + if stage.relaxed_fallback_used: + stage_name = f"{stage.name} (relaxed fallback)" + stage_table.add_row( + stage_name, + str(stage.count_in), + str(stage.count_out), + str(stage.dropped) if stage.dropped is not None else "-", + f"{stage.ms:.2f}" if stage.ms is not None else "-", + ) + console.print(stage_table) + + returned = [ + candidate for candidate in response.candidates if candidate.disposition == "returned" + ] + result_table = Table(title="Ranked results", show_header=True, header_style="bold") + result_table.add_column("Rank", justify="right") + result_table.add_column("Result") + result_table.add_column("Score", justify="right") + result_table.add_column("Δ", justify="right") + for candidate in returned: + label = _query_candidate_label(candidate) + result = Text(label) + if candidate.type == "entity" and candidate.permalink: + result.append(f"\n{candidate.permalink}", style="green") + if show_ids and (candidate_id := _query_candidate_id(candidate)) is not None: + result.append(f"\n{candidate_id}", style="dim") + result_table.add_row( + str(candidate.scores.final_rank or "-"), + result, + _score_text(candidate), + _movement_text(candidate), + ) + console.print(result_table) + + if not show_misses: + return + if response.retrieval_mode == SearchRetrievalMode.FTS: + console.print("[yellow]show-misses not applicable: FTS window is the SQL LIMIT[/yellow]") + return + + console.print("[dim]Misses: bounded window — not exhaustive[/dim]") + misses = sorted( + (candidate for candidate in response.candidates if candidate.disposition != "returned"), + key=lambda candidate: ( + candidate.disposition, + candidate.type or "", + candidate.id if candidate.id is not None else -1, + _query_candidate_label(candidate), + ), + ) + for disposition, grouped in groupby(misses, key=lambda candidate: candidate.disposition): + miss_table = Table(title=disposition, show_header=True, header_style="bold") + miss_table.add_column("Candidate") + miss_table.add_column("Score", justify="right") + miss_table.add_column("Detail") + for candidate in grouped: + detail = candidate.rejection_detail + detail_text = detail.model_dump_json(exclude_none=True) if detail is not None else "-" + miss_table.add_row( + _query_candidate_label(candidate), + ( + f"{candidate.scores.vector_similarity:.4f}" + if candidate.scores.vector_similarity is not None + else "-" + ), + detail_text, + ) + console.print(miss_table) + + +def _plain_query( + response: InspectQueryResponse, + *, + show_misses: bool, + show_ids: bool, +) -> None: + """Render the same query plan as undecorated, greppable text.""" + typer.echo(f"Retrieval query: {response.query}") + typer.echo(f"Mode: {response.retrieval_mode.value}") + typer.echo(f"Project: {response.project_id}") + typer.echo("Engine:") + typer.echo(f" Vector index: {response.engine.vector_index}") + typer.echo(f" Embedding model: {response.engine.embedding_model}") + if response.engine.ready_rows is None: + typer.echo(" Readiness: n/a (FTS execution reads no vector manifest)") + else: + typer.echo( + " Readiness: " + f"ready={response.engine.ready_rows} pending={response.engine.pending_rows} " + f"other_identity={response.engine.other_identity_rows}" + ) + typer.echo(f" Fusion: {response.engine.fusion_formula}") + reranker = response.engine.reranker + if reranker.enabled: + status = "applied" if reranker.applied else f"skipped={reranker.skipped_reason}" + typer.echo(f" Reranker: {reranker.model} candidates={reranker.candidates} {status}") + else: + typer.echo(" Reranker: disabled") + + typer.echo("Stages:") + for stage in response.stages: + ms = f"{stage.ms:.2f}" if stage.ms is not None else "-" + relaxed = " relaxed_fallback=yes" if stage.relaxed_fallback_used else "" + typer.echo( + f" {stage.name} in={stage.count_in} out={stage.count_out} " + f"dropped={stage.dropped if stage.dropped is not None else '-'} ms={ms}{relaxed}" + ) + + typer.echo("Ranked results:") + returned = [ + candidate for candidate in response.candidates if candidate.disposition == "returned" + ] + for candidate in returned: + label = _query_candidate_label(candidate) + identity = [] + if candidate.type == "entity" and candidate.permalink: + identity.append(f"permalink={candidate.permalink}") + if show_ids and (candidate_id := _query_candidate_id(candidate)) is not None: + identity.append(f"id={candidate_id}") + identity_text = f" {' '.join(identity)}" if identity else "" + typer.echo( + f" {candidate.scores.final_rank or '-'} {_score_text(candidate)} " + f"delta={_movement_text(candidate)} {label}{identity_text}" + ) + + if not show_misses: + return + if response.retrieval_mode == SearchRetrievalMode.FTS: + typer.echo("show-misses not applicable: FTS window is the SQL LIMIT") + return + typer.echo("Misses (bounded window — not exhaustive):") + misses = sorted( + (candidate for candidate in response.candidates if candidate.disposition != "returned"), + key=lambda candidate: ( + candidate.disposition, + candidate.type or "", + candidate.id if candidate.id is not None else -1, + _query_candidate_label(candidate), + ), + ) + for disposition, grouped in groupby(misses, key=lambda candidate: candidate.disposition): + typer.echo(f" {disposition}:") + for candidate in grouped: + typer.echo(f" {_query_candidate_label(candidate)}") + + +@inspect_app.command("query") +def inspect_query( + query_text: Annotated[str, typer.Argument(help="Text query to inspect")], + mode: InspectQueryMode = typer.Option( + InspectQueryMode.TEXT, + "--mode", + help="Retrieval mode: text, vector, or hybrid", + ), + show_misses: bool = typer.Option( + False, + "--show-misses", + help="Render bounded rejected candidates in human output", + ), + show_ids: bool = typer.Option( + False, + "--show-ids", + help="Include stable entity IDs in human output, with search-row fallbacks", + ), + page: int = typer.Option(1, "--page", min=1, help="Result page to inspect"), + page_size: int = typer.Option( + 10, + "--page-size", + min=1, + help="Results per page", + ), + json_output: bool = typer.Option(False, "--json", help="Output raw JSON"), + plain: bool = typer.Option(False, "--plain", help="Output undecorated plain text"), + project: Annotated[ + Optional[str], + typer.Option(help="The project to use; defaults to the configured project."), + ] = None, + project_id: Annotated[ + Optional[str], + typer.Option( + "--project-id", + help="Project external_id (UUID); takes precedence over --project.", + ), + ] = None, + local: bool = typer.Option( + False, "--local", help="Force local API routing (ignore cloud mode)" + ), + cloud: bool = typer.Option(False, "--cloud", help="Force cloud API routing"), +) -> None: + """Show the retrieval stages that produced one query result page.""" + from basic_memory.cli.commands.command_utils import run_with_cleanup + from fastmcp.exceptions import ToolError + + retrieval_mode = { + InspectQueryMode.TEXT: SearchRetrievalMode.FTS, + InspectQueryMode.VECTOR: SearchRetrievalMode.VECTOR, + InspectQueryMode.HYBRID: SearchRetrievalMode.HYBRID, + }[mode] + try: + validate_routing_flags(local, cloud) + _validate_output_flags(json_output, plain) + with force_routing(local=local, cloud=cloud): + response = run_with_cleanup( + run_inspect_query( + SearchQuery(text=query_text, retrieval_mode=retrieval_mode), + limit=page_size, + offset=(page - 1) * page_size, + project=project, + project_id=project_id, + ) + ) + + output_mode = _resolve_output_mode(json_output, plain) + if output_mode == "json": + print(response.model_dump_json(indent=2)) + elif output_mode == "plain": + _plain_query(response, show_misses=show_misses, show_ids=show_ids) + else: + _display_query(response, show_misses=show_misses, show_ids=show_ids) + except typer.Exit: + raise + except (ToolError, ValueError) as exc: + typer.echo(f"Error: {exc}", err=True) + raise typer.Exit(1) + except Exception as exc: # pragma: no cover + logger.error(f"Error inspecting retrieval query: {exc}") + typer.echo(f"Error: {exc}", err=True) + raise typer.Exit(1) + + @inspect_app.command("chunks") def inspect_chunks( identifier: Annotated[str, typer.Argument(help="Note identifier to inspect")], diff --git a/src/basic_memory/mcp/clients/inspect.py b/src/basic_memory/mcp/clients/inspect.py index b9e28ed88..c17454110 100644 --- a/src/basic_memory/mcp/clients/inspect.py +++ b/src/basic_memory/mcp/clients/inspect.py @@ -3,7 +3,13 @@ from httpx import AsyncClient import logfire -from basic_memory.schemas.inspect import InspectChunksRequest, InspectChunksResponse +from basic_memory.schemas.inspect import ( + InspectChunksRequest, + InspectChunksResponse, + InspectQueryRequest, + InspectQueryResponse, +) +from basic_memory.schemas.search import SearchQuery class InspectClient: @@ -32,3 +38,29 @@ async def inspect_chunks(self, identifier: str) -> InspectChunksResponse: path_template="/v2/projects/{project_id}/inspect/chunks", ) return InspectChunksResponse.model_validate(response.json()) + + async def inspect_query( + self, + query: SearchQuery, + *, + limit: int, + offset: int, + ) -> InspectQueryResponse: + """Run one search and return its execution-native retrieval trace.""" + from basic_memory.mcp.tools.utils import call_post + + request = InspectQueryRequest(query=query, limit=limit, offset=offset) + with logfire.span( + "mcp.client.inspect.query", + client_name="inspect", + operation="query", + ): + response = await call_post( + self.http_client, + f"{self._base_path}/query", + json=request.model_dump(mode="json"), + client_name="inspect", + operation="query", + path_template="/v2/projects/{project_id}/inspect/query", + ) + return InspectQueryResponse.model_validate(response.json()) diff --git a/src/basic_memory/repository/postgres_search_repository.py b/src/basic_memory/repository/postgres_search_repository.py index 6c1e7b810..ce3d405de 100644 --- a/src/basic_memory/repository/postgres_search_repository.py +++ b/src/basic_memory/repository/postgres_search_repository.py @@ -3,6 +3,7 @@ import asyncio import json import re +import time from collections.abc import Sequence from datetime import datetime from typing import Any, override, List, Optional @@ -25,6 +26,10 @@ SearchRepositoryBase, VectorChunkState, ) +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.semantic_errors import SemanticDependenciesMissingError from basic_memory.repository.semantic_vector_index import SemanticVectorIndex @@ -427,8 +432,15 @@ async def _run_vector_query( 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) + return await super()._run_vector_query( + session, + query_embedding, + candidate_limit, + trace=trace, + ) @override def _vector_prepare_window_size(self) -> int: @@ -864,6 +876,8 @@ async def search( offset: int = 0, allow_relaxed: bool = False, session: AsyncSession | None = None, + *, + trace: SearchTraceCollector | None = None, ) -> List[SearchIndexRow]: """Search across all indexed content using PostgreSQL tsvector.""" # --- Dispatch vector / hybrid modes (shared logic) --- @@ -881,6 +895,7 @@ async def search( min_similarity=min_similarity, limit=limit, offset=offset, + trace=trace, ) if dispatched is not None: return dispatched @@ -934,6 +949,7 @@ async def search( """ 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 @@ -951,6 +967,7 @@ async def execute_rows(active_session: AsyncSession, query_params: dict[str, Any 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: @@ -968,6 +985,7 @@ async def run_search(active_session: AsyncSession): # 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 " @@ -985,17 +1003,27 @@ async def run_search(active_session: AsyncSession): active_session, {**params, "text": relaxed}, ) - return rows + return rows, relaxed_fallback_used try: if session is not None: - rows = await run_search(session) + rows, relaxed_fallback_used = await run_search(session) else: async with db.scoped_session(self.session_maker) as owned_session: - rows = await run_search(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 @@ -1003,6 +1031,16 @@ async def run_search(active_session: AsyncSession): 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: diff --git a/src/basic_memory/repository/search_repository.py b/src/basic_memory/repository/search_repository.py index a25d0b538..fc394723e 100644 --- a/src/basic_memory/repository/search_repository.py +++ b/src/basic_memory/repository/search_repository.py @@ -19,6 +19,7 @@ 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 ChunkManifestRow +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, @@ -45,6 +46,15 @@ def configured_embedding_model(self) -> str: ... @property def configured_vector_index(self) -> str: ... + @property + def configured_min_similarity(self) -> float: ... + + @property + def configured_reranker_model(self) -> str | None: ... + + @property + def configured_reranker_candidates(self) -> int: ... + async def init_search_index(self) -> None: """Initialize the search index schema.""" ... @@ -75,6 +85,8 @@ async def search( offset: int = 0, allow_relaxed: bool = False, session: AsyncSession | None = None, + *, + trace: SearchTraceCollector | None = None, ) -> List[SearchIndexRow]: """Search across indexed content.""" ... diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 9e834c7be..08e11e25b 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -32,6 +32,20 @@ validate_rerank_scores, ) from basic_memory.repository.search_index_row import SearchIndexRow +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_chunking import ( SemanticSourceRow, VectorChunkRecord, @@ -74,6 +88,7 @@ # 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 @@ -234,6 +249,21 @@ def configured_vector_index(self) -> str: """Return the configured vector-index identity.""" return self._semantic_vector_index_name + @property + def configured_min_similarity(self) -> float: + """Return the configured vector similarity floor.""" + return self._semantic_min_similarity + + @property + def configured_reranker_model(self) -> str | None: + """Return the configured reranker identity without invoking it.""" + return self._rerank_provider.model_name if self._rerank_provider is not None else None + + @property + def configured_reranker_candidates(self) -> int: + """Return the fixed maximum reranker pool size.""" + return self._reranker_candidates + # ------------------------------------------------------------------ # Abstract methods — FTS and schema (backend-specific) # ------------------------------------------------------------------ @@ -282,6 +312,8 @@ async def search( limit: int = 10, offset: int = 0, allow_relaxed: bool = False, + *, + trace: SearchTraceCollector | None = None, ) -> List[SearchIndexRow]: """Search across all indexed content. @@ -341,8 +373,16 @@ async def _run_vector_query( 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 [] @@ -352,7 +392,14 @@ async def _run_vector_query( query_embedding, limit=candidate_limit, ) - return await self._hydrate_vector_matches(session, matches) + 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: @@ -360,13 +407,50 @@ async def _run_vector_query( query_embedding, limit=scan_limit, ) - hydrated = await self._hydrate_vector_matches(session, matches) + 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 ): - return hydrated[:candidate_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. @@ -380,6 +464,8 @@ 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: @@ -424,7 +510,7 @@ async def _hydrate_vector_matches( for row in result.mappings().all() } ) - return [ + hydrated = [ { "entity_id": match.key.entity_id, "chunk_key": match.key.chunk_key, @@ -434,6 +520,51 @@ async def _hydrate_vector_matches( 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, @@ -1750,6 +1881,7 @@ async def _dispatch_retrieval_mode( min_similarity: Optional[float] = None, limit: int, offset: int, + trace: SearchTraceCollector | None = None, ) -> Optional[List[SearchIndexRow]]: """Dispatch vector or hybrid retrieval if requested. @@ -1783,6 +1915,7 @@ async def _dispatch_retrieval_mode( min_similarity=min_similarity, limit=limit, offset=offset, + trace=trace, ) if mode == SearchRetrievalMode.HYBRID.value: if not can_use_vector: @@ -1803,6 +1936,7 @@ async def _dispatch_retrieval_mode( min_similarity=min_similarity, limit=limit, offset=offset, + trace=trace, ) # FTS mode: return None to let the subclass handle it @@ -1888,6 +2022,7 @@ async def _rerank_and_paginate( 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. @@ -1923,10 +2058,14 @@ async def _rerank_and_paginate( 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 scores = validate_rerank_scores( await self._rerank_provider.rerank(query_text, documents), len(pool), @@ -1939,8 +2078,26 @@ async def _rerank_and_paginate( pool=len(pool), model=self._rerank_provider.model_name, ) - demoted_tail = self._demote_tail(tail, floor=reranked[-1].score or 0.0) - return (reranked + demoted_tail)[offset:page_end] + 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, @@ -1960,6 +2117,7 @@ async def _search_vector_only( 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. @@ -1987,25 +2145,52 @@ async def _search_vector_only( # 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, - ) + 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) - vector_rows = await self._run_vector_query( - session, - query_embedding, - candidate_limit, - ) + 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 @@ -2070,6 +2255,16 @@ def _log_vector_summary() -> None: 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 } @@ -2082,6 +2277,15 @@ def _log_vector_summary() -> None: # 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( @@ -2115,6 +2319,13 @@ def _log_vector_summary() -> None: # Use (type, id) tuples to avoid collisions between different # search_index row types that share the same auto-increment id. allowed_keys = {(row.type, row.id) for row in filtered_rows if row.id is not None} + 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] = [] @@ -2152,6 +2363,8 @@ def _log_vector_summary() -> None: 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, @@ -2168,6 +2381,7 @@ def _log_vector_summary() -> None: candidate_limit=stable_candidate_limit, _emit_observability_log=False, _apply_rerank=False, + trace=None, ) output = await self._rerank_and_paginate( query_text, @@ -2175,6 +2389,7 @@ def _log_vector_summary() -> None: offset=offset, limit=limit, stable_rows=stable_rows, + trace=trace, ) else: output = ranked_rows[offset : offset + limit] @@ -2238,6 +2453,7 @@ async def _search_hybrid( _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. @@ -2273,6 +2489,7 @@ async def _search_hybrid( limit=candidate_limit, offset=0, allow_relaxed=True, + trace=trace, ) fts_ms = (time.perf_counter() - fts_start) * 1000 vector_start = time.perf_counter() @@ -2296,8 +2513,38 @@ async def _search_hybrid( 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() # --- Score-based fusion keyed on (type, id) --- @@ -2326,6 +2573,18 @@ async def _search_hybrid( 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, + 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): @@ -2349,6 +2608,18 @@ async def _search_hybrid( 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_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 @@ -2369,6 +2640,8 @@ def _materialize(entry: tuple[SearchIndexKey, float]) -> SearchIndexRow: 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, @@ -2385,6 +2658,7 @@ def _materialize(entry: tuple[SearchIndexKey, float]) -> SearchIndexRow: _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] @@ -2411,10 +2685,10 @@ def _materialize(entry: tuple[SearchIndexKey, float]) -> SearchIndexRow: offset=offset, limit=limit, stable_rows=stable_candidates, + trace=trace, ) else: output = [_materialize(entry) for entry in ranked[offset : offset + limit]] - fusion_ms = (time.perf_counter() - fusion_start) * 1000 total_ms = (time.perf_counter() - query_start) * 1000 if _emit_observability_log and total_ms > 2500: logger.warning( diff --git a/src/basic_memory/repository/search_trace.py b/src/basic_memory/repository/search_trace.py new file mode 100644 index 000000000..0540db9ac --- /dev/null +++ b/src/basic_memory/repository/search_trace.py @@ -0,0 +1,627 @@ +"""Typed, execution-native trace values for the search retrieval pipeline.""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any, Literal + +type TraceKey = tuple[str, int] +type RetrievalMode = Literal["fts", "vector", "hybrid"] +type DropReason = Literal[ + "not_in_manifest", + "pending", + "model_mismatch", + "index_mismatch", + "readiness_changed", + "malformed_key", +] +type MinSimilaritySource = Literal["query", "config"] + +HYDRATION_DROP_CLASSIFICATION_BATCH_SIZE = 250 + + +# --- Rejections --- + + +@dataclass(frozen=True, slots=True) +class BelowThreshold: + key: TraceKey + similarity: float + threshold: float + + +@dataclass(frozen=True, slots=True) +class HydrationDropped: + entity_id: int + chunk_key: str + similarity: float + reason: DropReason + stored_model: str | None + stored_index: str | None + + +@dataclass(frozen=True, slots=True) +class FilteredOut: + key: TraceKey + + +@dataclass(frozen=True, slots=True) +class MissingSearchRow: + key: TraceKey + + +@dataclass(frozen=True, slots=True) +class BeyondPageWindow: + key: TraceKey + rank: int + score: float + + +type Rejection = ( + BelowThreshold | HydrationDropped | FilteredOut | MissingSearchRow | BeyondPageWindow +) + + +# --- Frozen retrieval stages --- + + +@dataclass(frozen=True, slots=True) +class VectorChunkMatch: + key: TraceKey + chunk_key: str + similarity: float + # Owner entity of the search row this chunk serves — known at hydration and kept + # so rejected candidates can still be enriched with their stable external id. + entity_id: int | None + + +@dataclass(frozen=True, slots=True) +class VectorStageTrace: + candidate_limit: int + adapter_match_count: int + hydrated_count: int + drops: tuple[HydrationDropped, ...] + effective_min_similarity: float + min_similarity_source: MinSimilaritySource + threshold_rejections: tuple[BelowThreshold, ...] + filter_rejections: tuple[FilteredOut, ...] + missing_search_rows: tuple[MissingSearchRow, ...] + chunk_matches: tuple[VectorChunkMatch, ...] + embed_ms: float + vector_query_ms: float + + +@dataclass(frozen=True, slots=True) +class FtsScore: + key: TraceKey + score: float + rank: int + + +@dataclass(frozen=True, slots=True) +class FtsStageTrace: + raw_scores: tuple[FtsScore, ...] + normalized_scores: tuple[FtsScore, ...] | None + fts_max_abs: float | None + result_count: int + relaxed_fallback_used: bool + fts_ms: float | None + + +@dataclass(frozen=True, slots=True) +class FusionEntry: + key: TraceKey + fts_score: float | None + fts_rank: int | None + vector_score: float | None + vector_rank: int | None + fused_score: float + fused_rank: int + dual_source: bool + + +@dataclass(frozen=True, slots=True) +class FusionStageTrace: + formula_version: str + bonus: float + entries: tuple[FusionEntry, ...] + fusion_ms: float + + +@dataclass(frozen=True, slots=True) +class RerankEntry: + key: TraceKey + pre_rerank_rank: int + pre_rerank_score: float + rerank_score: float | None + post_rerank_rank: int + demoted_score: float | None + + +@dataclass(frozen=True, slots=True) +class RerankStageTrace: + provider_model: str + pool_size: int + reranker_candidates: int + entries: tuple[RerankEntry, ...] + tail_floor: float + stable_pool_refetched: bool + rerank_ms: float + + +@dataclass(frozen=True, slots=True) +class ManifestReadiness: + configured_index: str + configured_model: str + ready_rows: int + pending_rows: int + other_identity_rows: int + + +@dataclass(frozen=True, slots=True) +class FinalResultEntry: + key: TraceKey + entity_id: int | None + title: str | None + permalink: str | None + file_path: str + final_rank: int + final_score: float + + +@dataclass(frozen=True, slots=True) +class RerankerConfigSummary: + enabled: bool + model: str | None + candidates: int + + +@dataclass(frozen=True, slots=True) +class QueryMeta: + query_text: str + retrieval_mode: RetrievalMode + limit: int + offset: int + project_id: int + candidate_limit: int + rerank_pool_size: int + embedding_model: str + vector_index: str + fusion_formula_version: str + min_similarity: float + min_similarity_source: MinSimilaritySource + reranker: RerankerConfigSummary + rerank_applied: bool + rerank_skipped_reason: str | None + total_ms: float + + +@dataclass(frozen=True, slots=True) +class FtsQueryTrace: + meta: QueryMeta + fts: FtsStageTrace + final: tuple[FinalResultEntry, ...] + + +@dataclass(frozen=True, slots=True) +class VectorQueryTrace: + meta: QueryMeta + readiness: ManifestReadiness + vector: VectorStageTrace + rerank: RerankStageTrace | None + final: tuple[FinalResultEntry, ...] + + +@dataclass(frozen=True, slots=True) +class HybridQueryTrace: + meta: QueryMeta + readiness: ManifestReadiness + fts: FtsStageTrace + vector: VectorStageTrace + fusion: FusionStageTrace + rerank: RerankStageTrace | None + final: tuple[FinalResultEntry, ...] + + +type QueryTrace = FtsQueryTrace | VectorQueryTrace | HybridQueryTrace + + +@dataclass +class SearchTraceCollector: + """Mutable call-scoped accumulator, frozen into a mode-specific trace at the boundary.""" + + vector: VectorStageTrace | None = None + fts: FtsStageTrace | None = None + fusion: FusionStageTrace | None = None + rerank: RerankStageTrace | None = None + readiness: ManifestReadiness | None = None + stable_pool_refetched: bool = False + # Rendered from the exact prepared query the repository executed (including + # legacy note-type expansion), so the trace never re-derives its criteria. + executed_query_description: str | None = None + + +# --- Pure stage builders --- + + +def build_vector_stage( + *, + previous: VectorStageTrace | None = None, + candidate_limit: int | None = None, + adapter_match_count: int | None = None, + hydrated_count: int | None = None, + drops: Sequence[HydrationDropped] | None = None, + effective_min_similarity: float | None = None, + min_similarity_source: MinSimilaritySource | None = None, + threshold_rejections: Sequence[BelowThreshold] | None = None, + filter_rejections: Sequence[FilteredOut] | None = None, + missing_search_rows: Sequence[MissingSearchRow] | None = None, + chunk_matches: Mapping[TraceKey, Sequence[tuple[str, float, int | None]]] | None = None, + embed_ms: float | None = None, + vector_query_ms: float | None = None, +) -> VectorStageTrace: + """Freeze vector-stage values already captured by the executing pipeline.""" + if previous is None and ( + candidate_limit is None or adapter_match_count is None or hydrated_count is None + ): + raise ValueError("An initial vector stage requires candidate and adapter counts.") + + previous_candidate_limit = previous.candidate_limit if previous is not None else 0 + previous_adapter_match_count = previous.adapter_match_count if previous is not None else 0 + previous_hydrated_count = previous.hydrated_count if previous is not None else 0 + previous_chunk_matches: dict[TraceKey, list[tuple[str, float, int | None]]] = {} + if previous is not None: + for match in previous.chunk_matches: + previous_chunk_matches.setdefault(match.key, []).append( + (match.chunk_key, match.similarity, match.entity_id) + ) + effective_chunk_matches = chunk_matches if chunk_matches is not None else previous_chunk_matches + return VectorStageTrace( + candidate_limit=( + candidate_limit if candidate_limit is not None else previous_candidate_limit + ), + adapter_match_count=( + adapter_match_count if adapter_match_count is not None else previous_adapter_match_count + ), + hydrated_count=(hydrated_count if hydrated_count is not None else previous_hydrated_count), + drops=tuple(drops if drops is not None else (previous.drops if previous else ())), + effective_min_similarity=( + effective_min_similarity + if effective_min_similarity is not None + else (previous.effective_min_similarity if previous else 0.0) + ), + min_similarity_source=( + min_similarity_source + if min_similarity_source is not None + else (previous.min_similarity_source if previous else "config") + ), + threshold_rejections=tuple( + threshold_rejections + if threshold_rejections is not None + else (previous.threshold_rejections if previous else ()) + ), + filter_rejections=tuple( + filter_rejections + if filter_rejections is not None + else (previous.filter_rejections if previous else ()) + ), + missing_search_rows=tuple( + missing_search_rows + if missing_search_rows is not None + else (previous.missing_search_rows if previous else ()) + ), + chunk_matches=tuple( + VectorChunkMatch( + key=key, + chunk_key=chunk_key, + similarity=similarity, + entity_id=entity_id, + ) + for key, matches in effective_chunk_matches.items() + for chunk_key, similarity, entity_id in matches + ), + embed_ms=embed_ms if embed_ms is not None else (previous.embed_ms if previous else 0.0), + vector_query_ms=( + vector_query_ms + if vector_query_ms is not None + else (previous.vector_query_ms if previous else 0.0) + ), + ) + + +def build_fts_page_stage( + raw_scores: Sequence[tuple[TraceKey, float]], + *, + normalized_scores: Mapping[TraceKey, float] | None = None, + fts_max_abs: float | None = None, + relaxed_fallback_used: bool, + fts_ms: float | None = None, +) -> FtsStageTrace: + """Freeze the exact SQL page and optional hybrid normalization. + + ``raw_scores`` is the ordered page as executed — SQLite can serve duplicate + copies of one logical row, and each occurrence keeps its own rank here so the + stage never disagrees with the final results built from the same page. + """ + raw = tuple( + FtsScore(key=key, score=score, rank=rank) + for rank, (key, score) in enumerate(raw_scores, start=1) + ) + normalized = ( + tuple( + FtsScore(key=key, score=score, rank=rank) + for rank, (key, score) in enumerate(normalized_scores.items(), start=1) + ) + if normalized_scores is not None + else None + ) + return FtsStageTrace( + raw_scores=raw, + normalized_scores=normalized, + fts_max_abs=fts_max_abs, + result_count=len(raw), + relaxed_fallback_used=relaxed_fallback_used, + fts_ms=fts_ms, + ) + + +def build_fusion_stage( + *, + formula_version: str, + bonus: float, + fts_scores: Mapping[TraceKey, float], + fts_ranks: Mapping[TraceKey, int], + vector_scores: Mapping[TraceKey, float], + vector_ranks: Mapping[TraceKey, int], + ranked_scores: Sequence[tuple[TraceKey, float]], + fusion_ms: float, +) -> FusionStageTrace: + """Freeze each source leg and the exact fused ordering.""" + return FusionStageTrace( + formula_version=formula_version, + bonus=bonus, + entries=tuple( + FusionEntry( + key=key, + fts_score=fts_scores.get(key), + fts_rank=(fts_ranks[key] + 1 if key in fts_ranks else None), + vector_score=vector_scores.get(key), + vector_rank=(vector_ranks[key] + 1 if key in vector_ranks else None), + fused_score=score, + fused_rank=rank, + dual_source=key in fts_scores and key in vector_scores, + ) + for rank, (key, score) in enumerate(ranked_scores, start=1) + ), + fusion_ms=fusion_ms, + ) + + +def build_rerank_stage( + *, + provider_model: str, + reranker_candidates: int, + pre_rerank_scores: Mapping[TraceKey, float], + pool_keys: Sequence[TraceKey], + rerank_scores: Mapping[TraceKey, float], + post_rerank_rows: Sequence[tuple[TraceKey, float]], + demoted_scores: Mapping[TraceKey, float], + tail_floor: float, + stable_pool_refetched: bool, + rerank_ms: float, +) -> RerankStageTrace: + """Freeze pre-rewrite scores and the final pool-plus-demoted-tail ordering.""" + pre_ranks = {key: rank for rank, key in enumerate(pre_rerank_scores, start=1)} + return RerankStageTrace( + provider_model=provider_model, + pool_size=len(pool_keys), + reranker_candidates=reranker_candidates, + entries=tuple( + RerankEntry( + key=key, + pre_rerank_rank=pre_ranks[key], + pre_rerank_score=pre_rerank_scores[key], + rerank_score=rerank_scores.get(key), + post_rerank_rank=rank, + demoted_score=demoted_scores.get(key), + ) + for rank, (key, _score) in enumerate(post_rerank_rows, start=1) + ), + tail_floor=tail_floor, + stable_pool_refetched=stable_pool_refetched, + rerank_ms=rerank_ms, + ) + + +def finalize_query_trace( + collector: SearchTraceCollector, + meta: QueryMeta, + final: Sequence[FinalResultEntry], + mode: RetrievalMode, +) -> QueryTrace: + """Freeze the collector into the only valid stage combination for ``mode``.""" + final_entries = tuple(final) + match mode: + case "fts": + if collector.fts is None: + raise ValueError("Cannot finalize FTS query trace without an FTS stage.") + if any( + stage is not None + for stage in ( + collector.readiness, + collector.vector, + collector.fusion, + collector.rerank, + ) + ): + raise ValueError("FTS query trace contains incompatible semantic stages.") + return FtsQueryTrace(meta=meta, fts=collector.fts, final=final_entries) + case "vector": + if collector.readiness is None or collector.vector is None: + raise ValueError( + "Cannot finalize vector query trace without readiness and vector stages." + ) + if collector.fts is not None or collector.fusion is not None: + raise ValueError("Vector query trace contains incompatible FTS or fusion stages.") + return VectorQueryTrace( + meta=meta, + readiness=collector.readiness, + vector=collector.vector, + rerank=collector.rerank, + final=final_entries, + ) + case "hybrid": + if ( + collector.readiness is None + or collector.fts is None + or collector.vector is None + or collector.fusion is None + ): + raise ValueError( + "Cannot finalize hybrid query trace without readiness, FTS, vector, and " + "fusion stages." + ) + return HybridQueryTrace( + meta=meta, + readiness=collector.readiness, + fts=collector.fts, + vector=collector.vector, + fusion=collector.fusion, + rerank=collector.rerank, + final=final_entries, + ) + + +# --- Trace-only manifest reads --- + + +@dataclass(frozen=True, slots=True) +class HydrationDropKey: + entity_id: int + chunk_key: str + similarity: float + configured_index: str + configured_model: str + + +async def read_manifest_readiness( + session: Any, + project_id: int, + vector_index: str, + embedding_model: str, +) -> ManifestReadiness: + """Count configured readiness and rows stored under another vector identity.""" + from sqlalchemy import text + + readiness_result = await session.execute( + text( + "SELECT embedding_status, COUNT(*) AS row_count " + "FROM search_vector_chunks WHERE project_id = :project_id " + "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, + }, + ) + counts = { + str(row["embedding_status"]): int(row["row_count"]) + for row in readiness_result.mappings().all() + } + other_result = await session.execute( + text( + "SELECT COUNT(*) FROM search_vector_chunks WHERE project_id = :project_id " + "AND (vector_index <> :vector_index OR embedding_model <> :embedding_model)" + ), + { + "project_id": project_id, + "vector_index": vector_index, + "embedding_model": embedding_model, + }, + ) + return ManifestReadiness( + configured_index=vector_index, + configured_model=embedding_model, + ready_rows=counts.get("ready", 0), + pending_rows=counts.get("pending", 0), + other_identity_rows=int(other_result.scalar_one()), + ) + + +async def classify_hydration_drops( + session: Any, + project_id: int, + dropped_keys: Sequence[HydrationDropKey], +) -> tuple[HydrationDropped, ...]: + """Classify adapter hits rejected by authoritative manifest hydration.""" + if not dropped_keys: + return () + + from sqlalchemy import text + + stored_by_key: dict[tuple[int, str], Any] = {} + for batch_start in range( + 0, + len(dropped_keys), + 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} + predicates: list[str] = [] + for index, dropped in enumerate(batch): + params[f"entity_id_{index}"] = dropped.entity_id + params[f"chunk_key_{index}"] = dropped.chunk_key + predicates.append( + f"(entity_id = :entity_id_{index} AND chunk_key = :chunk_key_{index})" + ) + + # Constraint: unhealthy indexes can return thousands of dropped adapter hits. + # SQLite caps expression depth and both backends cap bind parameters, so classify + # in the same fixed-size windows used by authoritative manifest hydration. + 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 (" + + " OR ".join(predicates) + + ")" + ), + params, + ) + stored_by_key.update( + {(int(row["entity_id"]), str(row["chunk_key"])): row for row in result.mappings().all()} + ) + + classified: list[HydrationDropped] = [] + for dropped in dropped_keys: + stored = stored_by_key.get((dropped.entity_id, dropped.chunk_key)) + if stored is None: + reason: DropReason = "not_in_manifest" + stored_model = None + stored_index = None + else: + stored_model = str(stored["embedding_model"]) + stored_index = str(stored["vector_index"]) + if stored_model != dropped.configured_model: + reason = "model_mismatch" + elif stored_index != dropped.configured_index: + reason = "index_mismatch" + elif str(stored["embedding_status"]) == "pending": + reason = "pending" + else: + # Trigger: derived readiness changed between hydration and classification. + # Why: PostgreSQL statement snapshots and concurrent embedding sync can make + # the follow-up read observe a now-ready row that the executing query dropped. + # Outcome: report the race explicitly without locking or failing inspection. + reason = "readiness_changed" + classified.append( + HydrationDropped( + entity_id=dropped.entity_id, + chunk_key=dropped.chunk_key, + similarity=dropped.similarity, + reason=reason, + stored_model=stored_model, + stored_index=stored_index, + ) + ) + return tuple(classified) diff --git a/src/basic_memory/repository/sqlite_search_repository.py b/src/basic_memory/repository/sqlite_search_repository.py index 62dee4899..d44e678b6 100644 --- a/src/basic_memory/repository/sqlite_search_repository.py +++ b/src/basic_memory/repository/sqlite_search_repository.py @@ -2,6 +2,7 @@ import asyncio import re +import time from collections.abc import Sequence from contextlib import asynccontextmanager from datetime import datetime @@ -28,6 +29,10 @@ 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 SearchRepositoryBase +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.semantic_errors import SemanticDependenciesMissingError from basic_memory.repository.semantic_vector_index import SemanticVectorIndex @@ -598,8 +603,15 @@ async def _run_vector_query( 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) + return await super()._run_vector_query( + session, + query_embedding, + candidate_limit, + trace=trace, + ) @override async def _delete_entity_chunks( @@ -960,6 +972,8 @@ async def search( offset: int = 0, allow_relaxed: bool = False, session: AsyncSession | None = None, + *, + trace: SearchTraceCollector | None = None, ) -> List[SearchIndexRow]: """Search across all indexed content using SQLite FTS5. @@ -983,6 +997,7 @@ async def search( min_similarity=min_similarity, limit=limit, offset=offset, + trace=trace, ) if dispatched is not None: return dispatched @@ -1030,10 +1045,12 @@ async def search( """ 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 @@ -1044,6 +1061,7 @@ async def run_search(active_session: AsyncSession): # ranks multi-term matches first. relaxed = self._relaxed_fts_text(search_text) if allow_relaxed and not rows else None if relaxed and params.get("text"): + relaxed_fallback_used = True params["text"] = relaxed logger.debug( "Strict SQLite FTS returned 0 results; retrying relaxed FTS query " @@ -1058,19 +1076,29 @@ async def run_search(active_session: AsyncSession): ): result = await active_session.execute(text(sql), params) rows = result.fetchall() - return rows + return rows, relaxed_fallback_used try: if session is not None: - rows = await run_search(session) + rows, relaxed_fallback_used = await run_search(session) else: async with db.scoped_session(self.session_maker) as owned_session: - rows = await run_search(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 @@ -1078,6 +1106,16 @@ async def run_search(active_session: AsyncSession): 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: diff --git a/src/basic_memory/schemas/inspect.py b/src/basic_memory/schemas/inspect.py index d024e6ffb..cf6660515 100644 --- a/src/basic_memory/schemas/inspect.py +++ b/src/basic_memory/schemas/inspect.py @@ -1,9 +1,18 @@ """API schemas for note-level retrieval inspection.""" +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field from datetime import datetime -from typing import Literal, Self, assert_never +from typing import TYPE_CHECKING, Literal, Self, assert_never + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from basic_memory.schemas.search import SearchQuery, SearchRetrievalMode -from pydantic import BaseModel, ConfigDict, model_validator +if TYPE_CHECKING: + from basic_memory.repository.search_trace import QueryTrace, Rejection, TraceKey type ChunkStatus = Literal["ready", "pending", "stale", "orphaned"] type InspectFreshness = Literal[ @@ -144,3 +153,714 @@ def validate_freshness_detail(self) -> Self: if not detail_is_valid: raise ValueError(f"Invalid detail for freshness={self.freshness}") return self + + +# --- Query execution trace --- + + +class InspectQueryRequest(BaseModel): + """One search request and the bounded result window to explain.""" + + query: SearchQuery + limit: int = Field(default=10, ge=1) + offset: int = Field(default=0, ge=0) + + +class InspectQueryWindow(BaseModel): + limit: int + offset: int + candidate_limit: int + rerank_pool: int + + +class InspectQueryReranker(BaseModel): + enabled: bool + model: str | None + candidates: int + applied: bool + skipped_reason: str | None + + +class InspectQueryEngine(BaseModel): + embedding_model: str + vector_index: str + # None = not applicable: an FTS-only execution never reads the vector manifest, + # so zeros here would misreport a populated index as empty. + ready_rows: int | None + pending_rows: int | None + other_identity_rows: int | None + fusion_formula: str + min_similarity: float + min_similarity_source: Literal["query", "config"] + reranker: InspectQueryReranker + + +class InspectQueryStage(BaseModel): + name: str + count_in: int + count_out: int + dropped: int | None + ms: float | None + # Only the fts stage sets this: True means strict matching found nothing and the + # displayed results came from the relaxed OR/prefix retry. + relaxed_fallback_used: bool | None = None + + +class InspectMatchedChunk(BaseModel): + chunk_key: str + similarity: float + + +class InspectDroppedChunk(BaseModel): + chunk_key: str + similarity: float + reason: Literal[ + "not_in_manifest", + "pending", + "model_mismatch", + "index_mismatch", + "readiness_changed", + "malformed_key", + ] + stored_model: str | None + stored_index: str | None + + +class InspectQueryRejectionDetail(BaseModel): + reason: str + chunk_key: str | None = None + similarity: float | None = None + threshold: float | None = None + stored_model: str | None = None + stored_index: str | None = None + rank: int | None = None + score: float | None = None + + +class InspectQueryScores(BaseModel): + fts_raw: float | None = None + fts_normalized: float | None = None + fts_rank: int | None = None + vector_similarity: float | None = None + vector_rank: int | None = None + fused_score: float | None = None + fused_rank: int | None = None + dual_source: bool | None = None + pre_rerank_rank: int | None = None + pre_rerank_score: float | None = None + rerank_score: float | None = None + post_rerank_rank: int | None = None + final_rank: int | None = None + final_score: float | None = None + + +type InspectDisposition = Literal[ + "returned", + "below_threshold", + "not_in_manifest", + "pending", + "model_mismatch", + "index_mismatch", + "readiness_changed", + "malformed_key", + "filtered_out", + "missing_search_row", + "beyond_page_window", +] + + +class InspectQueryCandidate(BaseModel): + type: str | None + id: int | None + external_id: str | None + title: str | None + permalink: str | None + file_path: str | None + disposition: InspectDisposition + rejection_detail: InspectQueryRejectionDetail | None + matched_chunks: list[InspectMatchedChunk] + dropped_chunks: list[InspectDroppedChunk] + scores: InspectQueryScores + + +class InspectQueryTimings(BaseModel): + total: float + embedding: float | None + vector_query: float | None + fts: float | None + fusion: float | None + rerank: float | None + + +class InspectQueryResponse(BaseModel): + """Machine-readable flattening of one execution-native query trace.""" + + query: str + retrieval_mode: SearchRetrievalMode + project_id: int + window: InspectQueryWindow + engine: InspectQueryEngine + stages: list[InspectQueryStage] + candidates: list[InspectQueryCandidate] + timings_ms: InspectQueryTimings + + +@dataclass(frozen=True, slots=True) +class _ForeignOwnerDropKey: + """Identity for a parseable drop whose owner differs from the live row's owner. + + Manifest uniqueness includes entity_id, so one parseable chunk_key can exist under + two owners (a stale adapter hit surviving a search-row ID reuse); merging their + evidence would misattribute one entity's miss — or external id — to the other. + """ + + entity_id: int | None + key: TraceKey + + +@dataclass(frozen=True, slots=True) +class _MalformedDropKey: + """Identity for a dropped chunk whose key failed to parse, scoped to its owner. + + Two entities can serve identically malformed keys; the raw key alone would merge + them and misattribute one entity's drop evidence to the other. + """ + + entity_id: int | None + raw_chunk_key: str + + +@dataclass(slots=True) +class _CandidateTrace: + key: TraceKey | None + raw_chunk_key: str | None = None + entity_id: int | None = None + title: str | None = None + permalink: str | None = None + file_path: str | None = None + rejection: Rejection | None = None + matched_chunks: list[InspectMatchedChunk] = field(default_factory=list) + dropped_chunks: list[InspectDroppedChunk] = field(default_factory=list) + scores: InspectQueryScores = field(default_factory=InspectQueryScores) + + +def _best_similarity(current: float | None, observed: float) -> float: + """Track the best similarity without clamping negative cosine scores at zero.""" + return observed if current is None else max(current, observed) + + +def _parse_trace_chunk_key(chunk_key: str) -> TraceKey: + parts = chunk_key.split(":") + if len(parts) < 3: + raise ValueError(f"Invalid traced chunk key: {chunk_key!r}") + return parts[0], int(parts[1]) + + +def _rejection_disposition(rejection: Rejection) -> InspectDisposition: + from basic_memory.repository.search_trace import ( + BelowThreshold, + BeyondPageWindow, + FilteredOut, + HydrationDropped, + MissingSearchRow, + ) + + match rejection: + case BelowThreshold(): + return "below_threshold" + case HydrationDropped(reason=reason): + return reason + case FilteredOut(): + return "filtered_out" + case MissingSearchRow(): + return "missing_search_row" + case BeyondPageWindow(): + return "beyond_page_window" + case unexpected: # pragma: no cover - Rejection is a closed union + assert_never(unexpected) + + +def _rejection_detail(rejection: Rejection) -> InspectQueryRejectionDetail: + from basic_memory.repository.search_trace import ( + BelowThreshold, + BeyondPageWindow, + FilteredOut, + HydrationDropped, + MissingSearchRow, + ) + + match rejection: + case BelowThreshold(similarity=similarity, threshold=threshold): + return InspectQueryRejectionDetail( + reason="below_threshold", + similarity=similarity, + threshold=threshold, + ) + case HydrationDropped( + chunk_key=chunk_key, + similarity=similarity, + reason=reason, + stored_model=stored_model, + stored_index=stored_index, + ): + return InspectQueryRejectionDetail( + reason=reason, + chunk_key=chunk_key, + similarity=similarity, + stored_model=stored_model, + stored_index=stored_index, + ) + case FilteredOut(): + return InspectQueryRejectionDetail(reason="filtered_out") + case MissingSearchRow(): + return InspectQueryRejectionDetail(reason="missing_search_row") + case BeyondPageWindow(rank=rank, score=score): + return InspectQueryRejectionDetail( + reason="beyond_page_window", + rank=rank, + score=score, + ) + case unexpected: # pragma: no cover - Rejection is a closed union + assert_never(unexpected) + + +def query_trace_response( + trace: QueryTrace, + entity_external_id_lookup: Mapping[int, str] | None = None, +) -> InspectQueryResponse: + """Flatten a closed query trace without re-running or inferring retrieval stages.""" + from basic_memory.repository.search_trace import ( + BeyondPageWindow, + FtsQueryTrace, + HydrationDropped, + HybridQueryTrace, + VectorQueryTrace, + ) + + candidates: dict[TraceKey | _MalformedDropKey | _ForeignOwnerDropKey, _CandidateTrace] = {} + external_ids = entity_external_id_lookup or {} + + def candidate(key: TraceKey) -> _CandidateTrace: + return candidates.setdefault(key, _CandidateTrace(key=key)) + + fts = trace.fts if isinstance(trace, (FtsQueryTrace, HybridQueryTrace)) else None + vector = trace.vector if isinstance(trace, (VectorQueryTrace, HybridQueryTrace)) else None + fusion = trace.fusion if isinstance(trace, HybridQueryTrace) else None + rerank = trace.rerank if isinstance(trace, (VectorQueryTrace, HybridQueryTrace)) else None + readiness = trace.readiness if isinstance(trace, (VectorQueryTrace, HybridQueryTrace)) else None + + # Seed owners already known from the final results before any drop processing: + # a row returned via FTS alone has no vector chunk match to establish its owner, + # and foreign-owner drop detection needs that identity to keep a stale same-key + # hit from attaching to the returned row. + for result in trace.final: + if result.entity_id is not None: + candidate(result.key).entity_id = result.entity_id + + if fts is not None: + fts_rank_offset = trace.meta.offset if isinstance(trace, FtsQueryTrace) else 0 + for score in fts.raw_scores: + entry = candidate(score.key) + # Duplicate SQL-page occurrences share one logical candidate; the page is + # rank-ordered, so the first occurrence carries the row's best source rank. + if entry.scores.fts_rank is None: + entry.scores.fts_raw = score.score + entry.scores.fts_rank = score.rank + fts_rank_offset + for score in fts.normalized_scores or (): + candidate(score.key).scores.fts_normalized = score.score + + if vector is not None: + best_vector_scores: dict[TraceKey, float] = {} + for match in vector.chunk_matches: + entry = candidate(match.key) + if entry.entity_id is None: + entry.entity_id = match.entity_id + entry.matched_chunks.append( + InspectMatchedChunk( + chunk_key=match.chunk_key, + similarity=match.similarity, + ) + ) + best_vector_scores[match.key] = _best_similarity( + best_vector_scores.get(match.key), + match.similarity, + ) + surviving_vector_keys = set(best_vector_scores) + for rejection in vector.drops: + try: + key = _parse_trace_chunk_key(rejection.chunk_key) + except (TypeError, ValueError): + entry = candidates.setdefault( + _MalformedDropKey(rejection.entity_id, rejection.chunk_key), + _CandidateTrace(key=None, raw_chunk_key=rejection.chunk_key), + ) + entry.entity_id = rejection.entity_id + entry.rejection = rejection + entry.dropped_chunks.append( + InspectDroppedChunk( + chunk_key=rejection.chunk_key, + similarity=rejection.similarity, + reason=rejection.reason, + stored_model=rejection.stored_model, + stored_index=rejection.stored_index, + ) + ) + entry.scores.vector_similarity = _best_similarity( + entry.scores.vector_similarity, + rejection.similarity, + ) + continue + entry = candidate(key) + # Trigger: this drop names a different owner than the candidate already has. + # Why: the same parseable chunk_key can exist under two entity_ids; merging + # would overwrite the owner and misattribute drop evidence or external ids. + # Outcome: foreign-owner drops get their own candidate, scored and rejected + # from their drops alone (they have no hydrated chunks by construction). + if entry.entity_id is not None and rejection.entity_id != entry.entity_id: + foreign = candidates.setdefault( + _ForeignOwnerDropKey(rejection.entity_id, key), + _CandidateTrace(key=key), + ) + foreign.entity_id = rejection.entity_id + foreign.dropped_chunks.append( + InspectDroppedChunk( + chunk_key=rejection.chunk_key, + similarity=rejection.similarity, + reason=rejection.reason, + stored_model=rejection.stored_model, + stored_index=rejection.stored_index, + ) + ) + current_rejection = foreign.rejection + if current_rejection is None or ( + isinstance(current_rejection, HydrationDropped) + and ( + rejection.similarity > current_rejection.similarity + or ( + rejection.similarity == current_rejection.similarity + and rejection.chunk_key < current_rejection.chunk_key + ) + ) + ): + foreign.rejection = rejection + foreign.scores.vector_similarity = _best_similarity( + foreign.scores.vector_similarity, + rejection.similarity, + ) + continue + entry.entity_id = rejection.entity_id + entry.dropped_chunks.append( + InspectDroppedChunk( + chunk_key=rejection.chunk_key, + similarity=rejection.similarity, + reason=rejection.reason, + stored_model=rejection.stored_model, + stored_index=rejection.stored_index, + ) + ) + # Trigger: one row has both a ready chunk and a rejected sibling chunk. + # Why: retrieval ranks the row from its ready chunks, so promoting a sibling's + # drop to the row would make the trace disagree with the result execution. + # Outcome: retain the dropped chunk as evidence, but reject and score the row + # from drops only when no hydrated chunk for that row survived. + if key not in surviving_vector_keys: + current_rejection = entry.rejection + if current_rejection is None or ( + isinstance(current_rejection, HydrationDropped) + and ( + rejection.similarity > current_rejection.similarity + or ( + rejection.similarity == current_rejection.similarity + and rejection.chunk_key < current_rejection.chunk_key + ) + ) + ): + entry.rejection = rejection + best_vector_scores[key] = _best_similarity( + best_vector_scores.get(key), + rejection.similarity, + ) + for key, score in best_vector_scores.items(): + candidate(key).scores.vector_similarity = score + for rejection in vector.threshold_rejections: + candidate(rejection.key).rejection = rejection + for rejection in vector.filter_rejections: + candidate(rejection.key).rejection = rejection + for rejection in vector.missing_search_rows: + candidate(rejection.key).rejection = rejection + surviving_vector_scores = [ + (key, score) + for key, score in best_vector_scores.items() + if candidate(key).rejection is None + ] + surviving_vector_scores.sort(key=lambda item: item[1], reverse=True) + for rank, (key, _score) in enumerate(surviving_vector_scores, start=1): + candidate(key).scores.vector_rank = rank + + if fusion is not None: + for fused in fusion.entries: + entry = candidate(fused.key) + # Reaching fusion means at least one retrieval leg admitted the row. Any + # vector-only rejection is therefore chunk/leg evidence, not its disposition. + entry.rejection = None + entry.scores.fts_normalized = fused.fts_score + entry.scores.fts_rank = fused.fts_rank + entry.scores.vector_similarity = fused.vector_score + entry.scores.vector_rank = fused.vector_rank + entry.scores.fused_score = fused.fused_score + entry.scores.fused_rank = fused.fused_rank + entry.scores.dual_source = fused.dual_source + + if rerank is not None: + for reranked in rerank.entries: + entry = candidate(reranked.key) + entry.scores.pre_rerank_rank = reranked.pre_rerank_rank + entry.scores.pre_rerank_score = reranked.pre_rerank_score + entry.scores.rerank_score = reranked.rerank_score + entry.scores.post_rerank_rank = reranked.post_rerank_rank + + returned_keys = {result.key for result in trace.final} + ranked_for_page: list[tuple[TraceKey, int, float]] + if rerank is not None: + ranked_for_page = [ + ( + entry.key, + entry.post_rerank_rank, + entry.rerank_score + if entry.rerank_score is not None + else (entry.demoted_score or 0.0), + ) + for entry in rerank.entries + ] + elif fusion is not None: + ranked_for_page = [ + (entry.key, entry.fused_rank, entry.fused_score) for entry in fusion.entries + ] + elif vector is not None: + ranked_for_page = [ + (entry.key, entry.scores.vector_rank, entry.scores.vector_similarity) + for entry in candidates.values() + if entry.rejection is None + and entry.key is not None + and entry.scores.vector_rank is not None + and entry.scores.vector_similarity is not None + ] + else: + ranked_for_page = [] + + for key, rank, score in ranked_for_page: + entry = candidate(key) + if key not in returned_keys and entry.rejection is None: + entry.rejection = BeyondPageWindow(key=key, rank=rank, score=score) + + for result in trace.final: + entry = candidate(result.key) + entry.entity_id = result.entity_id + entry.title = result.title + entry.permalink = result.permalink + entry.file_path = result.file_path + entry.rejection = None + + stages: list[InspectQueryStage] = [] + fts_stage = ( + InspectQueryStage( + name="fts", + count_in=fts.result_count, + count_out=fts.result_count, + dropped=0, + ms=fts.fts_ms, + relaxed_fallback_used=fts.relaxed_fallback_used, + ) + if fts is not None + else None + ) + # The hybrid pipeline awaits the FTS leg before embedding and vector retrieval, + # so the displayed plan lists FTS first there — inverted order would misread the + # per-stage timings when diagnosing hybrid latency. + if fts_stage is not None and isinstance(trace, HybridQueryTrace): + stages.append(fts_stage) + if vector is not None: + served_chunk_count = len(vector.chunk_matches) + collapsed_row_count = len({match.key for match in vector.chunk_matches}) + # Threshold, note-type filters, and search-row hydration reject whole rows + # after the collapse; without a stage of their own the plan would show more + # collapsed rows than ranked results with the losses unexplained. + row_rejection_count = ( + len(vector.threshold_rejections) + + len(vector.filter_rejections) + + len(vector.missing_search_rows) + ) + stages.extend( + [ + InspectQueryStage( + name="embedding", + count_in=1, + count_out=1, + dropped=0, + ms=vector.embed_ms, + ), + InspectQueryStage( + name="vector", + count_in=vector.adapter_match_count, + count_out=vector.hydrated_count, + dropped=max(0, vector.adapter_match_count - vector.hydrated_count), + ms=vector.vector_query_ms, + ), + InspectQueryStage( + name="row_collapse", + count_in=served_chunk_count, + count_out=collapsed_row_count, + dropped=None, + ms=None, + ), + InspectQueryStage( + name="row_filters", + count_in=collapsed_row_count, + count_out=collapsed_row_count - row_rejection_count, + dropped=row_rejection_count, + ms=None, + ), + ] + ) + # Trigger: an external-index overfetch hydrated more ready chunks than the + # candidate window the search consumes. + # Why: those chunks were truncated, not dropped by hydration — folding them + # into the vector stage would misattribute healthy chunks as failures. + # Outcome: the truncation appears as its own stage between vector and + # row_collapse, keeping every count at a single scope. + if served_chunk_count < vector.hydrated_count: + stages.insert( + len(stages) - 2, + InspectQueryStage( + name="candidate_window", + count_in=vector.hydrated_count, + count_out=served_chunk_count, + dropped=vector.hydrated_count - served_chunk_count, + ms=None, + ), + ) + if fts_stage is not None and not isinstance(trace, HybridQueryTrace): + stages.append(fts_stage) + if fusion is not None: + stages.append( + InspectQueryStage( + name="fusion", + count_in=len(fusion.entries), + count_out=len(fusion.entries), + dropped=0, + ms=fusion.fusion_ms, + ) + ) + if rerank is not None: + stages.append( + InspectQueryStage( + name="rerank", + count_in=len(rerank.entries), + count_out=len(rerank.entries), + dropped=0, + ms=rerank.rerank_ms, + ) + ) + + # SQLite FTS can return duplicate copies of one logical row; the trace's final + # entries preserve each occurrence, so the response must too — one returned + # candidate per occurrence, sharing the key's stage data but carrying its own + # final rank and score. + returned_final_keys = {result.key for result in trace.final} + response_candidates = [ + InspectQueryCandidate( + type=result.key[0], + id=result.key[1], + external_id=( + external_ids.get(result.entity_id) if result.entity_id is not None else None + ), + title=result.title, + permalink=result.permalink, + file_path=result.file_path, + disposition="returned", + rejection_detail=None, + matched_chunks=candidates[result.key].matched_chunks, + dropped_chunks=candidates[result.key].dropped_chunks, + scores=candidates[result.key].scores.model_copy( + update={"final_rank": result.final_rank, "final_score": result.final_score} + ), + ) + for result in trace.final + ] + response_candidates += [ + InspectQueryCandidate( + type=entry.key[0] if entry.key is not None else None, + id=entry.key[1] if entry.key is not None else None, + external_id=( + external_ids.get(entry.entity_id) if entry.entity_id is not None else None + ), + title=entry.title, + permalink=entry.permalink, + file_path=entry.file_path, + disposition=( + "returned" if entry.rejection is None else _rejection_disposition(entry.rejection) + ), + rejection_detail=( + _rejection_detail(entry.rejection) if entry.rejection is not None else None + ), + matched_chunks=entry.matched_chunks, + dropped_chunks=entry.dropped_chunks, + scores=entry.scores, + ) + for entry in sorted( + # A rejected candidate is a miss even when it shares its row key with a + # returned result (foreign-owner drops); only the returned rows themselves + # (rejection cleared by the final loop) are excluded here. + ( + entry + for entry in candidates.values() + if entry.rejection is not None or entry.key not in returned_final_keys + ), + key=lambda item: ( + item.scores.post_rerank_rank or 1_000_000, + _rejection_disposition(item.rejection) if item.rejection is not None else "", + item.key is None, + item.key or ("", 0), + item.raw_chunk_key or "", + ), + ) + ] + + return InspectQueryResponse( + query=trace.meta.query_text, + retrieval_mode=SearchRetrievalMode(trace.meta.retrieval_mode), + project_id=trace.meta.project_id, + window=InspectQueryWindow( + limit=trace.meta.limit, + offset=trace.meta.offset, + candidate_limit=trace.meta.candidate_limit, + rerank_pool=trace.meta.rerank_pool_size, + ), + engine=InspectQueryEngine( + embedding_model=trace.meta.embedding_model, + vector_index=trace.meta.vector_index, + ready_rows=readiness.ready_rows if readiness is not None else None, + pending_rows=readiness.pending_rows if readiness is not None else None, + other_identity_rows=(readiness.other_identity_rows if readiness is not None else None), + fusion_formula=trace.meta.fusion_formula_version, + min_similarity=trace.meta.min_similarity, + min_similarity_source=trace.meta.min_similarity_source, + reranker=InspectQueryReranker( + enabled=trace.meta.reranker.enabled, + model=trace.meta.reranker.model, + candidates=trace.meta.reranker.candidates, + applied=trace.meta.rerank_applied, + skipped_reason=trace.meta.rerank_skipped_reason, + ), + ), + stages=stages, + candidates=response_candidates, + timings_ms=InspectQueryTimings( + total=trace.meta.total_ms, + embedding=vector.embed_ms if vector is not None else None, + vector_query=vector.vector_query_ms if vector is not None else None, + fts=fts.fts_ms if fts is not None else None, + fusion=fusion.fusion_ms if fusion is not None else None, + rerank=rerank.rerank_ms if rerank is not None else None, + ), + ) diff --git a/src/basic_memory/services/retrieval_inspect.py b/src/basic_memory/services/retrieval_inspect.py index 2318127c5..37db20c82 100644 --- a/src/basic_memory/services/retrieval_inspect.py +++ b/src/basic_memory/services/retrieval_inspect.py @@ -1,5 +1,6 @@ """Read-only inspection of one entity's retrieval projections.""" +import time from collections.abc import Mapping from dataclasses import dataclass from typing import Literal, assert_never @@ -16,16 +17,26 @@ 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_trace import ( + FinalResultEntry, + QueryMeta, + QueryTrace, + RerankerConfigSummary, + SearchTraceCollector, + finalize_query_trace, +) from basic_memory.repository.semantic_chunking import ( build_entity_fingerprint, build_vector_chunk_records, ) from basic_memory.schemas.inspect import ChunkStatus +from basic_memory.schemas.search import SearchQuery from basic_memory.services.exceptions import FileOperationError from basic_memory.services.file_service import FileService -from basic_memory.services.search_service import entity_embeddings_enabled +from basic_memory.services.search_service import SearchService, entity_embeddings_enabled @dataclass(frozen=True, slots=True) @@ -466,3 +477,106 @@ async def inspect_entity_chunks( if (row_type, row_id) not in search_rows_by_key ), ) + + +async def explain_query( + search_service: SearchService, + query: SearchQuery, + *, + limit: int, + offset: int, +) -> QueryTrace: + """Run one real search and freeze its execution-native retrieval trace.""" + collector = SearchTraceCollector() + started_at = time.perf_counter() + results = await search_service.search( + query, + limit=limit, + offset=offset, + trace=collector, + ) + total_ms = (time.perf_counter() - started_at) * 1000 + + repository = search_service.repository + mode = query.retrieval_mode.value + reranker_model = repository.configured_reranker_model + rerank_applied = collector.rerank is not None + if rerank_applied: + rerank_skipped_reason = None + elif mode == "fts": + rerank_skipped_reason = "fts_mode" + elif reranker_model is None: + rerank_skipped_reason = "disabled" + elif not results: + # Trigger: the final page is empty while earlier stages captured candidates. + # Why: hydrated chunks can still lose their whole row to the similarity + # threshold, filters, or a missing search row — only candidates that survive + # ranking can make an offset page "out of range"; anything less is a + # genuinely empty retrieval, even with hydrated chunks in the trace. + # Outcome: distinguish the empty page from an empty ranked candidate set. + if collector.fusion is not None: + ranked_candidates = len(collector.fusion.entries) + elif collector.vector is not None: + served_rows = {match.key for match in collector.vector.chunk_matches} + ranked_candidates = ( + len(served_rows) + - len(collector.vector.threshold_rejections) + - len(collector.vector.filter_rejections) + - len(collector.vector.missing_search_rows) + ) + else: + ranked_candidates = 0 + rerank_skipped_reason = "page_out_of_range" if ranked_candidates > 0 else "no_candidates" + else: + rerank_skipped_reason = "not_applied" + + candidate_limit = collector.vector.candidate_limit if collector.vector is not None else limit + effective_min_similarity = ( + query.min_similarity + if query.min_similarity is not None + else repository.configured_min_similarity + ) + meta = QueryMeta( + # Captured by search() from the exact prepared query it executed; absent only + # when preparation found no criteria and the search short-circuited. + query_text=( + collector.executed_query_description + if collector.executed_query_description is not None + else "(unsatisfiable query: no criteria)" + ), + retrieval_mode=mode, + limit=limit, + offset=offset, + project_id=repository.project_id, + candidate_limit=candidate_limit, + rerank_pool_size=(collector.rerank.pool_size if collector.rerank is not None else 0), + embedding_model=repository.configured_embedding_model, + vector_index=repository.configured_vector_index, + fusion_formula_version=FUSION_FORMULA_VERSION, + min_similarity=effective_min_similarity, + min_similarity_source=("query" if query.min_similarity is not None else "config"), + reranker=RerankerConfigSummary( + enabled=reranker_model is not None, + model=reranker_model, + candidates=repository.configured_reranker_candidates, + ), + rerank_applied=rerank_applied, + rerank_skipped_reason=rerank_skipped_reason, + total_ms=total_ms, + ) + final: list[FinalResultEntry] = [] + for rank, row in enumerate(results, start=1): + # entity_id is nullable storage on both backends (legacy/diagnostic rows); + # a missing owner degrades external-id enrichment, never the whole trace. + final.append( + FinalResultEntry( + key=(row.type, row.id), + entity_id=row.entity_id, + title=row.title, + permalink=row.permalink, + file_path=row.file_path, + final_rank=offset + rank, + final_score=row.score or 0.0, + ) + ) + return finalize_query_trace(collector, meta, final, mode) diff --git a/src/basic_memory/services/search_service.py b/src/basic_memory/services/search_service.py index 077b980b9..58de9cbaa 100644 --- a/src/basic_memory/services/search_service.py +++ b/src/basic_memory/services/search_service.py @@ -24,6 +24,7 @@ SearchRepository, ) from basic_memory.repository.search_query import relaxed_query_words +from basic_memory.repository.search_trace import SearchTraceCollector from basic_memory.schemas.base import normalize_note_type from basic_memory.schemas.search import SearchQuery, SearchItemType, SearchRetrievalMode from basic_memory.runtime.vector_sync import ( @@ -38,7 +39,7 @@ @dataclass(frozen=True) -class _PreparedSearchQuery: +class PreparedSearchQuery: """Normalized query inputs shared by search and count.""" search_text: str | None @@ -82,6 +83,36 @@ def entity_embeddings_enabled(entity: Entity) -> bool: return True +def describe_search_criteria(prepared: PreparedSearchQuery) -> str: + """Render the criteria the repository actually executed. + + The prepared query is the execution-native source: shorthand like ``text="tag:x"`` + normalizes into metadata filters, convenience fields fold into their canonical + forms, and note-type filters expand to legacy spellings, so describing the raw + request would misreport what ran. + """ + + def quoted(value: str | None) -> str | None: + return f'"{value}"' if value else None + + criteria: dict[str, object | None] = { + "text": quoted(prepared.search_text), + "title": quoted(prepared.title), + "permalink": quoted(prepared.permalink), + "permalink_match": quoted(prepared.permalink_match), + "note_types": list(prepared.note_types) if prepared.note_types else None, + # SearchItemType is a str-backed Enum whose str() is "SearchItemType.ENTITY"; + # the executed repository filter uses the plain value. + "entity_types": [item.value for item in prepared.search_item_types] + if prepared.search_item_types + else None, + "after_date": prepared.after_date, + "categories": list(prepared.categories) if prepared.categories else None, + "metadata_filters": dict(prepared.metadata_filters) if prepared.metadata_filters else None, + } + return " ".join(f"{name}={value}" for name, value in criteria.items() if value is not None) + + def _strip_nul(value: str) -> str: """Strip NUL bytes that PostgreSQL text columns cannot store. @@ -136,7 +167,7 @@ async def reindex_all(self, background_tasks: Optional[BackgroundTasks] = None) logger.info("Reindex complete") - def _prepare_query(self, query: SearchQuery) -> _PreparedSearchQuery | None: + def prepare_query(self, query: SearchQuery) -> PreparedSearchQuery | None: """Normalize a SearchQuery into repository arguments.""" search_text = query.text tags = query.tags @@ -170,7 +201,7 @@ def _prepare_query(self, query: SearchQuery) -> _PreparedSearchQuery | None: if query.status: metadata_filters.setdefault("status", query.status) - prepared = _PreparedSearchQuery( + prepared = PreparedSearchQuery( search_text=search_text, permalink=query.permalink, permalink_match=query.permalink_match, @@ -205,7 +236,7 @@ def _prepare_query(self, query: SearchQuery) -> _PreparedSearchQuery | None: return prepared @staticmethod - def _prepared_has_filters(prepared: _PreparedSearchQuery) -> bool: + def _prepared_has_filters(prepared: PreparedSearchQuery) -> bool: return bool( prepared.metadata_filters or prepared.note_types @@ -216,10 +247,10 @@ def _prepared_has_filters(prepared: _PreparedSearchQuery) -> bool: async def _include_legacy_note_type_spellings( self, - prepared: _PreparedSearchQuery, + prepared: PreparedSearchQuery, *, session: AsyncSession | None = None, - ) -> _PreparedSearchQuery: + ) -> PreparedSearchQuery: """Expand canonical note-type filters to exact legacy entity spellings.""" if not prepared.note_types: return prepared @@ -245,14 +276,33 @@ async def _include_legacy_note_type_spellings( async def _search_repository( self, - prepared: _PreparedSearchQuery, + prepared: PreparedSearchQuery, *, search_text: str | None, limit: int, offset: int, allow_relaxed: bool = False, session: AsyncSession | None = None, + trace: SearchTraceCollector | None = None, ) -> List[SearchIndexRow]: + if trace is None: + return await self.repository.search( + search_text=search_text, + permalink=prepared.permalink, + permalink_match=prepared.permalink_match, + title=prepared.title, + note_types=prepared.note_types, + search_item_types=prepared.search_item_types, + categories=prepared.categories, + after_date=prepared.after_date, + metadata_filters=prepared.metadata_filters, + retrieval_mode=prepared.retrieval_mode, + min_similarity=prepared.min_similarity, + limit=limit, + offset=offset, + allow_relaxed=allow_relaxed, + session=session, + ) return await self.repository.search( search_text=search_text, permalink=prepared.permalink, @@ -269,11 +319,12 @@ async def _search_repository( offset=offset, allow_relaxed=allow_relaxed, session=session, + trace=trace, ) async def _count_repository( self, - prepared: _PreparedSearchQuery, + prepared: PreparedSearchQuery, *, search_text: str | None, allow_relaxed: bool = False, @@ -299,6 +350,8 @@ async def search( limit=10, offset=0, session: AsyncSession | None = None, + *, + trace: SearchTraceCollector | None = None, ) -> List[SearchIndexRow]: """Search across all indexed content. @@ -307,13 +360,17 @@ async def search( 2. Pattern match: handles * wildcards in paths 3. Text search: full-text search across title/content """ - prepared = self._prepare_query(query) + prepared = self.prepare_query(query) if prepared is None: return [] prepared = await self._include_legacy_note_type_spellings( prepared, session=session, ) + if trace is not None: + # The trace must describe this execution's criteria, not a re-preparation: + # the legacy note-type expansion above depends on stored entity spellings. + trace.executed_query_description = describe_search_criteria(prepared) strict_search_text = prepared.search_text has_query = bool( @@ -344,13 +401,14 @@ async def search( offset=offset, allow_relaxed=allow_relaxed, session=session, + trace=trace, ) return results async def count(self, query: SearchQuery) -> int: """Count all indexed rows matching a query.""" - prepared = self._prepare_query(query) + prepared = self.prepare_query(query) if prepared is None: return 0 prepared = await self._include_legacy_note_type_spellings(prepared) diff --git a/tests/api/v2/test_inspect_router.py b/tests/api/v2/test_inspect_router.py index a92fe889d..f0ff7f3dd 100644 --- a/tests/api/v2/test_inspect_router.py +++ b/tests/api/v2/test_inspect_router.py @@ -1,18 +1,43 @@ """API contract tests for note-level retrieval inspection.""" +from importlib import import_module from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock import pytest from httpx import AsyncClient from sqlalchemy import text from basic_memory import db +from basic_memory.deps.services import get_search_service_v2_external from basic_memory.models import Project +from basic_memory.repository.semantic_errors import ( + RerankProviderContractError, + RerankTransientError, +) from basic_memory.repository.semantic_chunking import ( build_entity_fingerprint, build_vector_chunk_records, ) -from basic_memory.schemas.inspect import InspectChunksResponse, InspectRowsBehindFileDetail +from basic_memory.repository.search_trace import ( + FinalResultEntry, + HydrationDropped, + ManifestReadiness, + QueryMeta, + RerankerConfigSummary, + VectorQueryTrace, + build_vector_stage, +) +from basic_memory.schemas.inspect import ( + InspectChunksResponse, + InspectQueryRequest, + InspectQueryResponse, + InspectRowsBehindFileDetail, +) +from basic_memory.schemas.search import SearchQuery, SearchRetrievalMode + +inspect_router_module = import_module("basic_memory.api.v2.routers.inspect_router") async def _create_indexed_entity( @@ -250,3 +275,166 @@ async def test_inspect_chunks_semantic_disabled_returns_rows_only( assert inspection.detached == [] assert inspection.entity_fingerprint_indexed is None assert inspection.stale is False + + +@pytest.mark.asyncio +async def test_inspect_query_returns_schema_for_seeded_fts_corpus( + client: AsyncClient, + v2_project_url: str, + test_project: Project, + entity_repository, + search_service, + file_service, +): + entity = await _create_indexed_entity( + test_project=test_project, + title="API Query Inspection", + file_name="api-query-inspection.md", + entity_repository=entity_repository, + search_service=search_service, + file_service=file_service, + ) + + response = await client.post( + f"{v2_project_url}/inspect/query", + json={ + "query": {"text": "inspection", "retrieval_mode": "fts"}, + "limit": 5, + "offset": 0, + }, + ) + + assert response.status_code == 200, response.text + inspection = InspectQueryResponse.model_validate(response.json()) + assert inspection.query == 'text="inspection"' + assert inspection.retrieval_mode.value == "fts" + assert inspection.window.limit == 5 + assert inspection.candidates[0].permalink == entity.permalink + assert inspection.candidates[0].external_id == str(entity.external_id) + assert inspection.candidates[0].disposition == "returned" + assert [stage.name for stage in inspection.stages] == ["fts"] + + +@pytest.mark.asyncio +async def test_inspect_query_batch_enriches_final_and_hydration_drop_external_ids(monkeypatch): + trace = VectorQueryTrace( + meta=QueryMeta( + query_text="inspection", + retrieval_mode="vector", + limit=5, + offset=0, + project_id=1, + candidate_limit=10, + rerank_pool_size=0, + embedding_model="embedding", + vector_index="index", + fusion_formula_version="max+0.3*min/v1", + min_similarity=0.2, + min_similarity_source="config", + reranker=RerankerConfigSummary(enabled=False, model=None, candidates=20), + rerank_applied=False, + rerank_skipped_reason="disabled", + total_ms=1.0, + ), + readiness=ManifestReadiness("index", "embedding", 0, 0, 0), + vector=build_vector_stage( + candidate_limit=10, + adapter_match_count=1, + hydrated_count=0, + drops=( + HydrationDropped( + entity_id=2, + chunk_key="entity:2:0", + similarity=0.8, + reason="not_in_manifest", + stored_model=None, + stored_index=None, + ), + ), + ), + rerank=None, + final=( + FinalResultEntry( + key=("entity", 1), + entity_id=1, + title="One", + permalink="one", + file_path="one.md", + final_rank=1, + final_score=0.9, + ), + ), + ) + + async def fake_explain_query(*_args, **_kwargs): + return trace + + entity_service = MagicMock() + entity_service.get_entities_by_id = AsyncMock( + return_value=[ + SimpleNamespace(id=1, external_id="external-1"), + SimpleNamespace(id=2, external_id="external-2"), + ] + ) + monkeypatch.setattr(inspect_router_module, "explain_query", fake_explain_query) + + inspection = await inspect_router_module.inspect_query( + data=InspectQueryRequest( + query=SearchQuery(text="inspection", retrieval_mode=SearchRetrievalMode.VECTOR) + ), + project_id=1, + entity_service=entity_service, + search_service=MagicMock(), + ) + + assert {candidate.external_id for candidate in inspection.candidates} == { + "external-1", + "external-2", + } + entity_service.get_entities_by_id.assert_awaited_once_with([1, 2]) + + +@pytest.mark.asyncio +async def test_inspect_query_maps_semantic_disabled_to_400( + client: AsyncClient, + v2_project_url: str, +): + response = await client.post( + f"{v2_project_url}/inspect/query", + json={"query": {"text": "inspection", "retrieval_mode": "vector"}}, + ) + + assert response.status_code == 400 + assert "Semantic search is disabled" in response.json()["detail"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("error", "status_code"), + [ + (RerankTransientError("reranker unavailable"), 503), + (RerankProviderContractError("malformed reranker response"), 502), + ], +) +async def test_inspect_query_maps_reranker_errors( + client: AsyncClient, + v2_project_url: str, + app, + error: Exception, + status_code: int, +): + class RaisingSearchService: + async def search(self, *args, **kwargs): + raise error + + app.dependency_overrides[get_search_service_v2_external] = lambda: RaisingSearchService() + try: + response = await client.post( + f"{v2_project_url}/inspect/query", + json={"query": {"text": "inspection", "retrieval_mode": "hybrid"}}, + ) + finally: + app.dependency_overrides.pop(get_search_service_v2_external, None) + + assert response.status_code == status_code + assert response.json()["detail"] == str(error) diff --git a/tests/cli/test_db_reindex.py b/tests/cli/test_db_reindex.py index 6dfd2e8b7..7f73ed778 100644 --- a/tests/cli/test_db_reindex.py +++ b/tests/cli/test_db_reindex.py @@ -823,9 +823,7 @@ def test_reindex_total_embedding_failure_surfaces_error_and_exits_one( "[cyan]model=FastEmbedEmbeddingProvider:BAAI/bge-small-en-v1.5[/cyan]): " "0 entities embedded, 0 skipped, 3 errors" ) in output - representative_line = next( - line for line in printed_lines if "Representative error:" in line - ) + representative_line = next(line for line in printed_lines if "Representative error:" in line) assert "first line second line" in representative_line assert representative_line.endswith("...") assert "Reindex failed: all vector embedding attempts failed." in output diff --git a/tests/cli/test_inspect_command.py b/tests/cli/test_inspect_command.py index d85f1335d..54a263536 100644 --- a/tests/cli/test_inspect_command.py +++ b/tests/cli/test_inspect_command.py @@ -20,9 +20,20 @@ InspectDetachedSearchRow, InspectFreshness, InspectIndexBehindRowsDetail, + InspectMatchedChunk, + InspectQueryCandidate, + InspectQueryEngine, + InspectQueryRejectionDetail, + InspectQueryResponse, + InspectQueryReranker, + InspectQueryScores, + InspectQueryStage, + InspectQueryTimings, + InspectQueryWindow, InspectRowsBehindFileDetail, InspectSearchRow, ) +from basic_memory.schemas.search import SearchQuery, SearchRetrievalMode runner = CliRunner() @@ -160,6 +171,104 @@ def _response_with_freshness(freshness: InspectFreshness) -> InspectChunksRespon ) +def _query_response( + retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.HYBRID, +) -> InspectQueryResponse: + return InspectQueryResponse( + query="auth retrieval", + retrieval_mode=retrieval_mode, + project_id=7, + window=InspectQueryWindow( + limit=10, + offset=0, + candidate_limit=100, + rerank_pool=2, + ), + engine=InspectQueryEngine( + embedding_model="TraceEmbedding:model:4", + vector_index="sqlite-vec", + ready_rows=12, + pending_rows=2, + other_identity_rows=3, + fusion_formula="max+0.3*min/v1", + min_similarity=0.2, + min_similarity_source="config", + reranker=InspectQueryReranker( + enabled=True, + model="trace-reranker", + candidates=20, + applied=True, + skipped_reason=None, + ), + ), + stages=[ + InspectQueryStage( + name="fts", + count_in=3, + count_out=3, + dropped=0, + ms=0.8, + relaxed_fallback_used=True, + ), + InspectQueryStage(name="embedding", count_in=1, count_out=1, dropped=0, ms=1.2), + InspectQueryStage(name="vector", count_in=5, count_out=2, dropped=3, ms=2.3), + InspectQueryStage(name="fusion", count_in=3, count_out=3, dropped=0, ms=0.4), + InspectQueryStage(name="rerank", count_in=2, count_out=2, dropped=0, ms=4.5), + ], + candidates=[ + InspectQueryCandidate( + type="entity", + id=1, + external_id="11111111-1111-1111-1111-111111111111", + title="Authentication Guide", + permalink="notes/authentication-guide", + file_path="Authentication Guide.md", + disposition="returned", + rejection_detail=None, + matched_chunks=[InspectMatchedChunk(chunk_key="entity:1:0", similarity=0.9)], + dropped_chunks=[], + scores=InspectQueryScores( + vector_similarity=0.9, + vector_rank=2, + fused_score=1.1, + fused_rank=2, + pre_rerank_rank=2, + pre_rerank_score=1.1, + rerank_score=0.95, + post_rerank_rank=1, + final_rank=1, + final_score=0.95, + ), + ), + InspectQueryCandidate( + type="entity", + id=2, + external_id=None, + title=None, + permalink=None, + file_path=None, + disposition="below_threshold", + rejection_detail=InspectQueryRejectionDetail( + reason="below_threshold", + similarity=0.1, + threshold=0.2, + ), + matched_chunks=[InspectMatchedChunk(chunk_key="entity:2:0", similarity=0.1)], + dropped_chunks=[], + scores=InspectQueryScores(vector_similarity=0.1, vector_rank=5), + ), + ], + timings_ms=InspectQueryTimings( + total=8.4, + embedding=1.2, + vector_query=2.3, + fts=0.5, + fusion=0.4, + rerank=4.5, + ), + ) + + @patch("basic_memory.cli.commands.tool._use_rich", return_value=True) @patch("basic_memory.cli.commands.inspect.run_inspect_chunks", new_callable=AsyncMock) def test_inspect_chunks_rich_rendering(mock_run, _mock_use_rich): @@ -366,3 +475,233 @@ async def fake_get_project_client(*, project=None, project_id=None): assert await_args.args[1] == ( "/v2/projects/33333333-3333-3333-3333-333333333333/inspect/chunks" ) + + +@patch("basic_memory.cli.commands.tool._use_rich", return_value=True) +@patch("basic_memory.cli.commands.inspect.run_inspect_query", new_callable=AsyncMock) +def test_inspect_query_rich_renders_query_plan_and_grouped_misses(mock_run, _mock_use_rich): + mock_run.return_value = _query_response() + + result = runner.invoke( + cli_app, + [ + "inspect", + "query", + "auth retrieval", + "--mode", + "hybrid", + "--show-misses", + "--show-ids", + ], + ) + + assert result.exit_code == 0, result.output + assert "Retrieval query" in result.output + assert "Project: 7" in result.output + assert "12 ready, 2 pending, 3 other identity" in result.output + assert "max+0.3*min/v1" in result.output + assert "Ranked results" in result.output + assert "fts (relaxed fallback)" in result.output + assert "Authentication Guide" in result.output + assert "notes/authentication-guide" in result.output + assert "11111111-1111-1111-1111-111111111111" in result.output + assert "+1" in result.output + assert "bounded window — not exhaustive" in result.output + assert "below_threshold" in result.output + + +@patch("basic_memory.cli.commands.tool._use_rich", return_value=True) +@patch("basic_memory.cli.commands.inspect.run_inspect_query", new_callable=AsyncMock) +def test_inspect_query_renders_fts_readiness_as_not_applicable(mock_run, _mock_use_rich): + """An FTS-only trace has no manifest snapshot; zeros would misreport the index.""" + response = _query_response() + mock_run.return_value = response.model_copy( + update={ + "engine": response.engine.model_copy( + update={ + "ready_rows": None, + "pending_rows": None, + "other_identity_rows": None, + } + ) + } + ) + + rich_result = runner.invoke(cli_app, ["inspect", "query", "auth retrieval"]) + assert rich_result.exit_code == 0, rich_result.output + assert "Readiness: n/a" in rich_result.output + assert "0 ready" not in rich_result.output + + plain_result = runner.invoke(cli_app, ["inspect", "query", "auth retrieval", "--plain"]) + assert plain_result.exit_code == 0, plain_result.output + assert "Readiness: n/a" in plain_result.output + assert "ready=0" not in plain_result.output + + +@patch("basic_memory.cli.commands.inspect.run_inspect_query", new_callable=AsyncMock) +def test_inspect_query_plain_hides_misses_without_show_misses(mock_run): + mock_run.return_value = _query_response() + + result = runner.invoke( + cli_app, + ["inspect", "query", "auth retrieval", "--mode", "vector", "--plain"], + ) + + assert result.exit_code == 0, result.output + assert "Engine:" in result.output + assert "relaxed_fallback=yes" in result.output + assert "dropped=3" in result.output + assert "delta=+1" in result.output + assert "permalink=notes/authentication-guide" in result.output + assert "11111111-1111-1111-1111-111111111111" not in result.output + assert "id=entity:1" not in result.output + assert "bounded window" not in result.output + query = mock_run.await_args.args[0] + assert query.retrieval_mode == SearchRetrievalMode.VECTOR + + +@patch("basic_memory.cli.commands.inspect.run_inspect_query", new_callable=AsyncMock) +def test_inspect_query_json_is_schema_locked_and_always_includes_misses(mock_run): + expected = _query_response() + mock_run.return_value = expected + + result = runner.invoke( + cli_app, + [ + "inspect", + "query", + "auth retrieval", + "--json", + "--page", + "2", + "--page-size", + "5", + "--show-ids", + ], + ) + + assert result.exit_code == 0, result.output + validated = InspectQueryResponse.model_validate_json(result.output) + assert validated == expected + assert {candidate.disposition for candidate in validated.candidates} == { + "returned", + "below_threshold", + } + assert mock_run.await_args.kwargs["limit"] == 5 + assert mock_run.await_args.kwargs["offset"] == 5 + + +@patch("basic_memory.cli.commands.inspect.run_inspect_query", new_callable=AsyncMock) +def test_inspect_query_plain_show_ids_adds_external_id(mock_run): + mock_run.return_value = _query_response() + + result = runner.invoke( + cli_app, + ["inspect", "query", "auth retrieval", "--plain", "--show-ids"], + ) + + assert result.exit_code == 0, result.output + assert "permalink=notes/authentication-guide" in result.output + assert "id=11111111-1111-1111-1111-111111111111" in result.output + assert "id=entity:1" not in result.output + + +@patch("basic_memory.cli.commands.inspect.run_inspect_query", new_callable=AsyncMock) +def test_inspect_query_plain_show_ids_falls_back_to_type_qualified_trace_key(mock_run): + response = _query_response() + returned = response.candidates[0].model_copy(update={"external_id": None}) + mock_run.return_value = response.model_copy( + update={"candidates": [returned, *response.candidates[1:]]} + ) + + result = runner.invoke( + cli_app, + ["inspect", "query", "auth retrieval", "--plain", "--show-ids"], + ) + + assert result.exit_code == 0, result.output + assert "id=entity:1" in result.output + + +@patch("basic_memory.cli.commands.inspect.run_inspect_query", new_callable=AsyncMock) +def test_inspect_query_plain_preserves_malformed_drop_chunk_key(mock_run): + response = _query_response() + malformed = InspectQueryCandidate( + type=None, + id=None, + external_id=None, + title=None, + permalink=None, + file_path=None, + disposition="not_in_manifest", + rejection_detail=InspectQueryRejectionDetail( + reason="not_in_manifest", + chunk_key="malformed-key", + similarity=0.4, + ), + matched_chunks=[InspectMatchedChunk(chunk_key="malformed-key", similarity=0.4)], + dropped_chunks=[], + scores=InspectQueryScores(vector_similarity=0.4), + ) + mock_run.return_value = response.model_copy( + update={"candidates": [*response.candidates, malformed]} + ) + + result = runner.invoke( + cli_app, + ["inspect", "query", "auth retrieval", "--plain", "--show-misses"], + ) + + assert result.exit_code == 0, result.output + assert "malformed-key" in result.output + assert "None:None" not in result.output + + +@patch("basic_memory.cli.commands.inspect.run_inspect_query", new_callable=AsyncMock) +def test_inspect_query_fts_show_misses_explains_sql_limit(mock_run): + response = _query_response(SearchRetrievalMode.FTS) + response.engine.reranker.applied = False + response.engine.reranker.skipped_reason = "fts_mode" + mock_run.return_value = response + + result = runner.invoke( + cli_app, + ["inspect", "query", "auth retrieval", "--show-misses", "--plain"], + ) + + assert result.exit_code == 0, result.output + assert "show-misses not applicable: FTS window is the SQL LIMIT" in result.output + + +@pytest.mark.asyncio +async def test_run_inspect_query_uses_typed_client_and_project_route(monkeypatch): + expected = _query_response() + http_client = MagicMock() + active_project = SimpleNamespace(external_id="33333333-3333-3333-3333-333333333333") + + @asynccontextmanager + async def fake_get_project_client(*, project=None, project_id=None): + assert project == "research" + assert project_id == "33333333-3333-3333-3333-333333333333" + yield http_client, active_project + + response = MagicMock() + response.json.return_value = expected.model_dump(mode="json") + call_post = AsyncMock(return_value=response) + monkeypatch.setattr(inspect_command, "get_project_client", fake_get_project_client) + monkeypatch.setattr("basic_memory.mcp.tools.utils.call_post", call_post) + + query = SearchQuery(text="auth retrieval", retrieval_mode=SearchRetrievalMode.HYBRID) + result = await inspect_command.run_inspect_query( + query, + limit=10, + offset=0, + project="research", + project_id="33333333-3333-3333-3333-333333333333", + ) + + assert result == expected + call_post.assert_awaited_once() + await_args = call_post.await_args + assert await_args is not None + assert await_args.args[1] == ("/v2/projects/33333333-3333-3333-3333-333333333333/inspect/query") diff --git a/tests/repository/test_hybrid_fusion.py b/tests/repository/test_hybrid_fusion.py index 0b849974b..2042814ec 100644 --- a/tests/repository/test_hybrid_fusion.py +++ b/tests/repository/test_hybrid_fusion.py @@ -16,6 +16,7 @@ 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 FUSION_BONUS, SearchRepositoryBase +from basic_memory.repository.search_trace import SearchTraceCollector from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode @@ -85,6 +86,8 @@ async def search( limit: int = 10, offset: int = 0, allow_relaxed: bool = False, + *, + trace: SearchTraceCollector | None = None, ) -> list[SearchIndexRow]: return [] # pragma: no cover @@ -93,7 +96,14 @@ async def _ensure_vector_tables(self): pass # pragma: no cover @override - async def _run_vector_query(self, session, query_embedding, candidate_limit): + async def _run_vector_query( + self, + session, + query_embedding, + candidate_limit, + *, + trace: SearchTraceCollector | None = None, + ): return [] # pragma: no cover @override diff --git a/tests/repository/test_search_trace.py b/tests/repository/test_search_trace.py new file mode 100644 index 000000000..68c1dd3d2 --- /dev/null +++ b/tests/repository/test_search_trace.py @@ -0,0 +1,1386 @@ +"""Execution-native search trace builders and repository integration.""" + +from collections.abc import Sequence +from dataclasses import replace +from datetime import datetime, timezone +from typing import Any +from unittest.mock import MagicMock + +import pytest +from sqlalchemy import event, text + +from basic_memory import db +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_trace import ( + BelowThreshold, + FilteredOut, + FinalResultEntry, + FtsQueryTrace, + HydrationDropped, + HydrationDropKey, + HybridQueryTrace, + ManifestReadiness, + MissingSearchRow, + QueryMeta, + RetrievalMode, + RerankerConfigSummary, + SearchTraceCollector, + VectorQueryTrace, + build_fts_page_stage, + build_fusion_stage, + build_rerank_stage, + build_vector_stage, + classify_hydration_drops, + finalize_query_trace, +) +from basic_memory.repository.semantic_vector_index import ( + VectorDeletion, + VectorIndexScope, + VectorKey, + VectorMatch, + VectorRecord, +) +from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository +from basic_memory.schemas.inspect import query_trace_response +from basic_memory.schemas.search import SearchItemType, SearchQuery, SearchRetrievalMode +from basic_memory.services.retrieval_inspect import explain_query +from basic_memory.services.search_service import SearchService + +type BackendRepository = SQLiteSearchRepository | PostgresSearchRepository + + +class _TraceEmbeddingProvider: + model_name = "trace-embedding" + 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 _text in texts] + + def runtime_log_attrs(self) -> dict[str, Any]: + return {} + + +class _TraceVectorIndex: + def __init__(self, project_id: int) -> None: + self.matches: list[VectorMatch] = [] + self.scope = VectorIndexScope( + namespace="trace-test", + project_id=project_id, + embedding_identity="trace-embedding", + dimensions=4, + ) + + async def initialize(self) -> None: + return None + + async def upsert(self, records: Sequence[VectorRecord]) -> None: + return None + + async def delete(self, records: Sequence[VectorDeletion]) -> None: + return None + + async def delete_entity(self, entity_id: int) -> None: + return None + + async def search(self, query: Sequence[float], *, limit: int) -> list[VectorMatch]: + return self.matches[:limit] + + +class _ReverseReranker: + model_name = "reverse-reranker" + + async def rerank(self, query: str, documents: list[str]) -> list[float]: + return [0.1, 0.9][: len(documents)] + + def runtime_log_attrs(self) -> dict[str, Any]: + return {} + + +def _meta(mode: RetrievalMode) -> QueryMeta: + return QueryMeta( + query_text="alpha", + retrieval_mode=mode, + limit=2, + offset=0, + project_id=1, + candidate_limit=10, + rerank_pool_size=0, + embedding_model="embedding", + vector_index="index", + fusion_formula_version="max+0.3*min/v1", + min_similarity=0.2, + min_similarity_source="config", + reranker=RerankerConfigSummary(enabled=False, model=None, candidates=20), + rerank_applied=False, + rerank_skipped_reason="disabled", + total_ms=1.0, + ) + + +def test_stage_builders_freeze_plain_values_and_exact_fusion_math(): + vector = build_vector_stage( + candidate_limit=10, + adapter_match_count=3, + hydrated_count=2, + drops=( + HydrationDropped( + entity_id=3, + chunk_key="entity:3:0", + similarity=0.7, + reason="pending", + stored_model="embedding", + stored_index="index", + ), + ), + effective_min_similarity=0.5, + min_similarity_source="query", + threshold_rejections=(BelowThreshold(key=("entity", 2), similarity=0.4, threshold=0.5),), + filter_rejections=(FilteredOut(key=("entity", 4)),), + missing_search_rows=(MissingSearchRow(key=("entity", 5)),), + chunk_matches={ + ("entity", 1): [("entity:1:0", 0.8, 1)], + ("entity", 2): [("entity:2:0", 0.4, 2)], + }, + embed_ms=1.25, + vector_query_ms=2.5, + ) + assert vector.candidate_limit == 10 + assert vector.chunk_matches[0].chunk_key == "entity:1:0" + assert vector.threshold_rejections[0].threshold == 0.5 + + fts = build_fts_page_stage( + [(("entity", 1), -5.0), (("entity", 2), -2.5)], + normalized_scores={("entity", 1): 1.0, ("entity", 2): 0.5}, + fts_max_abs=5.0, + relaxed_fallback_used=True, + fts_ms=3.0, + ) + assert fts.relaxed_fallback_used is True + assert [entry.rank for entry in fts.raw_scores] == [1, 2] + + fusion = build_fusion_stage( + formula_version="max+0.3*min/v1", + bonus=FUSION_BONUS, + fts_scores={("entity", 1): 1.0}, + fts_ranks={("entity", 1): 0}, + vector_scores={("entity", 1): 0.8, ("entity", 2): 0.7}, + vector_ranks={("entity", 1): 0, ("entity", 2): 1}, + ranked_scores=[ + (("entity", 1), 1.0 + FUSION_BONUS * 0.8), + (("entity", 2), 0.7), + ], + fusion_ms=0.5, + ) + assert fusion.entries[0].fused_score == pytest.approx(1.24) + assert fusion.entries[0].dual_source is True + + rerank = build_rerank_stage( + provider_model="reranker", + reranker_candidates=2, + pre_rerank_scores={("entity", 1): 1.24, ("entity", 2): 0.7}, + pool_keys=[("entity", 1)], + rerank_scores={("entity", 1): 0.2}, + post_rerank_rows=[(("entity", 1), 0.2), (("entity", 2), 0.1)], + demoted_scores={("entity", 2): 0.1}, + tail_floor=0.2, + stable_pool_refetched=True, + rerank_ms=4.0, + ) + assert rerank.entries[0].pre_rerank_score == pytest.approx(1.24) + assert rerank.entries[1].rerank_score is None + assert rerank.entries[1].demoted_score == pytest.approx(0.1) + + +def test_fts_degenerate_stage_and_mode_specific_finalize_fail_fast(): + with pytest.raises(ValueError, match="initial vector stage"): + build_vector_stage() + + degenerate = build_fts_page_stage([], relaxed_fallback_used=True) + collector = SearchTraceCollector(fts=degenerate) + trace = finalize_query_trace(collector, _meta("fts"), (), "fts") + assert isinstance(trace, FtsQueryTrace) + assert trace.fts.result_count == 0 + assert trace.fts.relaxed_fallback_used is True + + with pytest.raises(ValueError, match="without readiness and vector stages"): + finalize_query_trace(SearchTraceCollector(), _meta("vector"), (), "vector") + with pytest.raises(ValueError, match="without readiness, FTS, vector, and fusion"): + finalize_query_trace(SearchTraceCollector(), _meta("hybrid"), (), "hybrid") + with pytest.raises(ValueError, match="without an FTS stage"): + finalize_query_trace(SearchTraceCollector(), _meta("fts"), (), "fts") + with pytest.raises(ValueError, match="incompatible semantic stages"): + finalize_query_trace( + SearchTraceCollector(fts=degenerate, readiness=ManifestReadiness("i", "m", 0, 0, 0)), + _meta("fts"), + (), + "fts", + ) + with pytest.raises(ValueError, match="incompatible FTS or fusion stages"): + finalize_query_trace( + SearchTraceCollector( + fts=degenerate, + vector=build_vector_stage( + candidate_limit=1, + adapter_match_count=0, + hydrated_count=0, + ), + readiness=ManifestReadiness("i", "m", 0, 0, 0), + ), + _meta("vector"), + (), + "vector", + ) + + +def test_query_response_flattens_every_rejection_variant(): + vector = build_vector_stage( + candidate_limit=10, + adapter_match_count=7, + hydrated_count=4, + drops=( + HydrationDropped(2, "entity:2:0", 0.8, "not_in_manifest", None, None), + HydrationDropped(3, "entity:3:0", 0.7, "pending", "m", "i"), + HydrationDropped(4, "entity:4:0", 0.6, "model_mismatch", "old", "i"), + HydrationDropped(5, "entity:5:0", 0.5, "index_mismatch", "m", "old"), + HydrationDropped(10, "entity:10:0", 0.45, "readiness_changed", "m", "i"), + HydrationDropped(11, "malformed-key", 0.35, "not_in_manifest", None, None), + ), + threshold_rejections=(BelowThreshold(("entity", 6), 0.4, 0.5),), + filter_rejections=(FilteredOut(("entity", 7)),), + missing_search_rows=(MissingSearchRow(("entity", 8)),), + chunk_matches={ + ("entity", 1): [("entity:1:0", 0.9, 1)], + ("entity", 6): [("entity:6:0", 0.4, 6)], + ("entity", 7): [("entity:7:0", 0.8, 7)], + ("entity", 8): [("entity:8:0", 0.7, 8)], + ("entity", 9): [("entity:9:0", 0.65, 9)], + }, + ) + collector = SearchTraceCollector( + vector=vector, + readiness=ManifestReadiness("i", "m", 4, 1, 2), + ) + final = (FinalResultEntry(("entity", 1), 1, "One", "one", "one.md", 1, 0.9),) + trace = finalize_query_trace(collector, _meta("vector"), final, "vector") + response = query_trace_response( + trace, + {1: "external-1", 2: "external-2", 6: "external-6", 11: "external-11"}, + ) + assert {candidate.disposition for candidate in response.candidates} == { + "returned", + "not_in_manifest", + "pending", + "model_mismatch", + "index_mismatch", + "readiness_changed", + "below_threshold", + "filtered_out", + "missing_search_row", + "beyond_page_window", + } + surviving = {candidate.id: candidate for candidate in response.candidates} + assert surviving[1].scores.vector_rank == 1 + assert surviving[1].external_id == "external-1" + assert surviving[2].external_id == "external-2" + # A hydrated match rejected by the threshold keeps its owner from the chunk match, + # so the response can still enrich it with a stable external id. + assert surviving[6].external_id == "external-6" + assert surviving[9].scores.vector_rank == 2 + vector_stage = next(stage for stage in response.stages if stage.name == "vector") + assert vector_stage.count_out == 4 + assert vector_stage.dropped == 3 + assert next(stage for stage in response.stages if stage.name == "row_collapse").dropped is None + # Threshold (entity 6), filter (entity 7), and missing-row (entity 8) rejections + # land after the collapse; the plan must show those losses, not jump from 5 to 1. + row_filters = next(stage for stage in response.stages if stage.name == "row_filters") + assert (row_filters.count_in, row_filters.count_out, row_filters.dropped) == (5, 2, 3) + malformed = next(candidate for candidate in response.candidates if candidate.id is None) + assert malformed.type is None + assert malformed.external_id == "external-11" + assert malformed.rejection_detail is not None + assert malformed.rejection_detail.chunk_key == "malformed-key" + assert malformed.matched_chunks == [] + assert [chunk.chunk_key for chunk in malformed.dropped_chunks] == ["malformed-key"] + + +def test_query_response_keeps_row_when_a_ready_chunk_survives_a_dropped_sibling(): + vector = build_vector_stage( + candidate_limit=10, + adapter_match_count=4, + hydrated_count=3, + drops=( + HydrationDropped( + 1, + "entity:1:1", + 0.95, + "pending", + "m", + "i", + ), + ), + chunk_matches={ + ("entity", 1): [("entity:1:0", 0.8, 1), ("entity:1:2", 0.7, 1)], + ("entity", 2): [("entity:2:0", 0.9, 2)], + }, + ) + trace = finalize_query_trace( + SearchTraceCollector( + vector=vector, + readiness=ManifestReadiness("i", "m", 2, 1, 0), + ), + replace(_meta("vector"), limit=1), + (FinalResultEntry(("entity", 2), 2, "Two", "two", "two.md", 1, 0.9),), + "vector", + ) + + response = query_trace_response(trace) + candidates = {candidate.id: candidate for candidate in response.candidates} + assert candidates[1].disposition == "beyond_page_window" + assert candidates[1].scores.vector_similarity == pytest.approx(0.8) + assert candidates[1].scores.vector_rank == 2 + assert [chunk.chunk_key for chunk in candidates[1].matched_chunks] == [ + "entity:1:0", + "entity:1:2", + ] + assert [(chunk.chunk_key, chunk.reason) for chunk in candidates[1].dropped_chunks] == [ + ("entity:1:1", "pending") + ] + vector_stage = next(stage for stage in response.stages if stage.name == "vector") + assert (vector_stage.count_in, vector_stage.count_out, vector_stage.dropped) == (4, 3, 1) + row_collapse = next(stage for stage in response.stages if stage.name == "row_collapse") + assert (row_collapse.count_in, row_collapse.count_out, row_collapse.dropped) == (3, 2, None) + row_filters = next(stage for stage in response.stages if stage.name == "row_filters") + assert (row_filters.count_in, row_filters.count_out, row_filters.dropped) == (2, 2, 0) + + +def test_query_response_uses_best_dropped_chunk_rejection_deterministically(): + vector = build_vector_stage( + candidate_limit=10, + adapter_match_count=3, + hydrated_count=0, + drops=( + HydrationDropped(1, "entity:1:2", 0.8, "pending", "m", "i"), + HydrationDropped(1, "entity:1:1", 0.9, "model_mismatch", "old", "i"), + HydrationDropped(1, "entity:1:0", 0.9, "index_mismatch", "m", "old"), + ), + ) + trace = finalize_query_trace( + SearchTraceCollector( + vector=vector, + readiness=ManifestReadiness("i", "m", 0, 3, 0), + ), + _meta("vector"), + (), + "vector", + ) + + candidate = query_trace_response(trace).candidates[0] + + assert candidate.disposition == "index_mismatch" + assert candidate.rejection_detail is not None + assert candidate.rejection_detail.chunk_key == "entity:1:0" + assert candidate.scores.vector_similarity == pytest.approx(0.9) + assert [chunk.chunk_key for chunk in candidate.dropped_chunks] == [ + "entity:1:2", + "entity:1:1", + "entity:1:0", + ] + + +def test_query_response_clears_vector_rejection_for_fts_fusion_survivor(): + fts = build_fts_page_stage( + [(("entity", 2), -1.0), (("entity", 1), -0.8)], + normalized_scores={("entity", 2): 1.0, ("entity", 1): 0.8}, + fts_max_abs=1.0, + relaxed_fallback_used=False, + ) + vector = build_vector_stage( + candidate_limit=10, + adapter_match_count=1, + hydrated_count=1, + threshold_rejections=(BelowThreshold(("entity", 1), 0.1, 0.5),), + chunk_matches={("entity", 1): [("entity:1:0", 0.1, 1)]}, + ) + fusion = build_fusion_stage( + formula_version="max+0.3*min/v1", + bonus=FUSION_BONUS, + fts_scores={("entity", 2): 1.0, ("entity", 1): 0.8}, + fts_ranks={("entity", 2): 0, ("entity", 1): 1}, + vector_scores={}, + vector_ranks={}, + ranked_scores=[(("entity", 2), 1.0), (("entity", 1), 0.8)], + fusion_ms=0.1, + ) + trace = finalize_query_trace( + SearchTraceCollector( + fts=fts, + vector=vector, + fusion=fusion, + readiness=ManifestReadiness("i", "m", 1, 0, 0), + ), + replace(_meta("hybrid"), limit=1), + (FinalResultEntry(("entity", 2), 2, "Two", "two", "two.md", 1, 1.0),), + "hybrid", + ) + + response = query_trace_response(trace) + candidates = {candidate.id: candidate for candidate in response.candidates} + + # Hybrid awaits FTS before embedding/vector, and the displayed plan follows suit. + assert [stage.name for stage in response.stages] == [ + "fts", + "embedding", + "vector", + "row_collapse", + "row_filters", + "fusion", + ] + assert candidates[1].disposition == "beyond_page_window" + assert candidates[1].rejection_detail is not None + assert candidates[1].rejection_detail.reason == "beyond_page_window" + + +def test_fts_query_response_uses_global_ranks_after_sql_offset(): + fts = build_fts_page_stage( + [(("entity", 11), -1.0), (("entity", 12), -0.8)], + relaxed_fallback_used=False, + ) + trace = finalize_query_trace( + SearchTraceCollector(fts=fts), + replace(_meta("fts"), limit=2, offset=10), + ( + FinalResultEntry(("entity", 11), 11, "Eleven", "eleven", "11.md", 11, -1.0), + FinalResultEntry(("entity", 12), 12, "Twelve", "twelve", "12.md", 12, -0.8), + ), + "fts", + ) + + response = query_trace_response(trace) + + assert [candidate.scores.fts_rank for candidate in response.candidates] == [11, 12] + # FTS execution never reads the vector manifest; zeros would misreport a + # populated index as empty. + assert response.engine.ready_rows is None + assert response.engine.pending_rows is None + assert response.engine.other_identity_rows is None + + +def _repository( + session_maker, + test_project, + app_config: BasicMemoryConfig, +) -> tuple[BackendRepository, _TraceVectorIndex]: + config = app_config.model_copy( + update={ + "semantic_search_enabled": True, + "semantic_min_similarity": 0.0, + "semantic_vector_k": 10, + } + ) + vector_index = _TraceVectorIndex(test_project.id) + repository_type = ( + PostgresSearchRepository + if config.database_backend == DatabaseBackend.POSTGRES + else SQLiteSearchRepository + ) + repository = repository_type( + session_maker, + project_id=test_project.id, + app_config=config, + embedding_provider=_TraceEmbeddingProvider(), + vector_index_name="trace-test", + vector_index=vector_index, + ) + return repository, vector_index + + +def _row(project_id: int, row_id: int, title: str, note_type: str) -> SearchIndexRow: + now = datetime.now(timezone.utc) + return SearchIndexRow( + project_id=project_id, + id=row_id, + type="entity", + title=title, + permalink=f"notes/{row_id}", + file_path=f"{row_id}.md", + metadata={"note_type": note_type}, + entity_id=row_id, + content_stems=f"{title} auth retrieval", + content_snippet=f"{title} auth retrieval", + created_at=now, + updated_at=now, + ) + + +async def _seed_trace_corpus( + repository: BackendRepository, + vector_index: _TraceVectorIndex, +) -> None: + await repository.init_search_index() + await repository.bulk_index_items( + [ + _row(repository.project_id, 1, "Alpha", "keep"), + _row(repository.project_id, 2, "Bravo", "keep"), + _row(repository.project_id, 3, "Charlie", "drop"), + ] + ) + configured_model = repository._embedding_model_key() + rows = [ + (1, "entity:1:0", "Alpha auth retrieval", configured_model, "trace-test", "ready"), + (2, "entity:2:0", "Bravo auth retrieval", configured_model, "trace-test", "ready"), + (3, "entity:3:0", "Charlie auth retrieval", configured_model, "trace-test", "ready"), + (999, "entity:999:0", "Missing search row", configured_model, "trace-test", "ready"), + (4, "entity:4:0", "Pending", configured_model, "trace-test", "pending"), + (5, "entity:5:0", "Old model", "old-model", "trace-test", "ready"), + (7, "entity:7:0", "Old index", configured_model, "old-index", "ready"), + ] + async with db.scoped_session(repository.session_maker) as session: + await session.execute( + text( + "INSERT INTO search_vector_chunks (" + "project_id, entity_id, chunk_key, chunk_text, source_hash, " + "entity_fingerprint, embedding_model, vector_index, embedding_status" + ") VALUES (" + ":project_id, :entity_id, :chunk_key, :chunk_text, 'hash', 'fingerprint', " + ":embedding_model, :vector_index, :embedding_status)" + ), + [ + { + "project_id": repository.project_id, + "entity_id": entity_id, + "chunk_key": chunk_key, + "chunk_text": chunk_text, + "embedding_model": embedding_model, + "vector_index": stored_index, + "embedding_status": status, + } + for entity_id, chunk_key, chunk_text, embedding_model, stored_index, status in rows + ], + ) + await session.commit() + + vector_index.matches = [ + VectorMatch(VectorKey(1, "entity:1:0"), 0.95), + VectorMatch(VectorKey(2, "entity:2:0"), 0.45), + VectorMatch(VectorKey(3, "entity:3:0"), 0.85), + VectorMatch(VectorKey(999, "entity:999:0"), 0.8), + VectorMatch(VectorKey(4, "entity:4:0"), 0.75), + VectorMatch(VectorKey(5, "entity:5:0"), 0.7), + VectorMatch(VectorKey(6, "entity:6:0"), 0.65), + VectorMatch(VectorKey(7, "entity:7:0"), 0.6), + ] + + +@pytest.mark.asyncio +async def test_vector_trace_captures_drops_threshold_filter_and_missing_on_same_execution( + session_maker, + test_project, + app_config, + engine_factory, +): + repository, vector_index = _repository(session_maker, test_project, app_config) + await _seed_trace_corpus(repository, vector_index) + + statements: list[str] = [] + engine, _session_maker = engine_factory + + def record_statement(_conn, _cursor, statement, _parameters, _context, _executemany): + statements.append(statement) + + event.listen(engine.sync_engine, "before_cursor_execute", record_statement) + try: + without_trace = await repository.search( + search_text="auth", + note_types=["keep"], + retrieval_mode=SearchRetrievalMode.VECTOR, + min_similarity=0.5, + limit=10, + ) + finally: + event.remove(engine.sync_engine, "before_cursor_execute", record_statement) + assert not any("GROUP BY embedding_status" in statement for statement in statements) + assert not any( + "embedding_model, vector_index, embedding_status" in statement for statement in statements + ) + + collector = SearchTraceCollector() + with_trace = await repository.search( + search_text="auth", + note_types=["keep"], + retrieval_mode=SearchRetrievalMode.VECTOR, + min_similarity=0.5, + limit=10, + trace=collector, + ) + assert [(row.type, row.id, row.score) for row in with_trace] == [ + (row.type, row.id, row.score) for row in without_trace + ] + assert [row.id for row in with_trace] == [1] + assert collector.readiness == ManifestReadiness( + configured_index="trace-test", + configured_model=repository._embedding_model_key(), + ready_rows=4, + pending_rows=1, + other_identity_rows=2, + ) + assert collector.vector is not None + assert {drop.reason for drop in collector.vector.drops} == { + "pending", + "model_mismatch", + "not_in_manifest", + "index_mismatch", + } + assert [rejection.key for rejection in collector.vector.threshold_rejections] == [("entity", 2)] + assert [rejection.key for rejection in collector.vector.filter_rejections] == [("entity", 3)] + assert [rejection.key for rejection in collector.vector.missing_search_rows] == [ + ("entity", 999) + ] + 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, ()) == () + readiness_race = await classify_hydration_drops( + session, + repository.project_id, + ( + HydrationDropKey( + entity_id=1, + chunk_key="entity:1:0", + similarity=0.95, + configured_index="trace-test", + configured_model=repository._embedding_model_key(), + ), + ), + ) + assert readiness_race[0].reason == "readiness_changed" + assert readiness_race[0].entity_id == 1 + + +@pytest.mark.asyncio +async def test_classify_hydration_drops_batches_large_unhealthy_candidate_set( + session_maker, + test_project, +): + dropped_keys = tuple( + HydrationDropKey( + entity_id=entity_id, + chunk_key=f"entity:{entity_id}:0", + similarity=0.5, + configured_index="trace-test", + configured_model="trace-embedding", + ) + for entity_id in range(10_000, 11_001) + ) + + async with db.scoped_session(session_maker) as session: + classified = await classify_hydration_drops( + session, + test_project.id, + dropped_keys, + ) + + assert len(classified) == len(dropped_keys) + assert {drop.reason for drop in classified} == {"not_in_manifest"} + assert [classified[0].entity_id, classified[-1].entity_id] == [10_000, 11_000] + + +@pytest.mark.asyncio +async def test_classify_hydration_drop_observes_pending_to_ready_transition( + session_maker, + test_project, + app_config, +): + repository, vector_index = _repository(session_maker, test_project, app_config) + await _seed_trace_corpus(repository, vector_index) + match = VectorMatch(VectorKey(1, "entity:1:0"), 0.95) + async with db.scoped_session(session_maker) as session: + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'pending' " + "WHERE project_id = :project_id AND entity_id = 1" + ), + {"project_id": test_project.id}, + ) + await session.commit() + async with db.scoped_session(session_maker) as session: + assert await repository._hydrate_vector_matches(session, [match]) == [] + + async with db.scoped_session(session_maker) as session: + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'ready' " + "WHERE project_id = :project_id AND entity_id = 1" + ), + {"project_id": test_project.id}, + ) + await session.commit() + async with db.scoped_session(session_maker) as session: + classified = await classify_hydration_drops( + session, + test_project.id, + ( + HydrationDropKey( + entity_id=1, + chunk_key="entity:1:0", + similarity=0.95, + configured_index="trace-test", + configured_model=repository._embedding_model_key(), + ), + ), + ) + + assert classified[0].reason == "readiness_changed" + + +@pytest.mark.asyncio +async def test_fts_vector_hybrid_and_rerank_trace_variants( + session_maker, + test_project, + app_config, +): + repository, vector_index = _repository(session_maker, test_project, app_config) + await _seed_trace_corpus(repository, vector_index) + search_service = SearchService(repository, MagicMock(), MagicMock(), session_maker) + + fts_trace = await explain_query( + search_service, + SearchQuery(text="Alpha", retrieval_mode=SearchRetrievalMode.FTS), + limit=10, + offset=0, + ) + assert isinstance(fts_trace, FtsQueryTrace) + assert fts_trace.fts.normalized_scores is None + assert fts_trace.fts.fts_ms is not None + assert fts_trace.fts.fts_ms >= 0.0 + assert fts_trace.final[0].key == ("entity", 1) + + relaxed_fts_trace = await explain_query( + search_service, + SearchQuery(text="Alpha absent-term", retrieval_mode=SearchRetrievalMode.FTS), + limit=10, + offset=0, + ) + assert isinstance(relaxed_fts_trace, FtsQueryTrace) + assert relaxed_fts_trace.fts.relaxed_fallback_used is True + assert relaxed_fts_trace.fts.fts_ms is not None + assert relaxed_fts_trace.fts.fts_ms >= 0.0 + + vector_trace = await explain_query( + search_service, + SearchQuery(text="auth", retrieval_mode=SearchRetrievalMode.VECTOR), + limit=2, + offset=0, + ) + assert isinstance(vector_trace, VectorQueryTrace) + assert vector_trace.vector.adapter_match_count == len(vector_index.matches) + + hybrid_trace = await explain_query( + search_service, + SearchQuery(text="Alpha", retrieval_mode=SearchRetrievalMode.HYBRID), + limit=3, + offset=0, + ) + assert isinstance(hybrid_trace, HybridQueryTrace) + alpha = next(entry for entry in hybrid_trace.fusion.entries if entry.key == ("entity", 1)) + assert alpha.fts_score == pytest.approx(1.0) + assert alpha.vector_score == pytest.approx(0.95) + assert alpha.fused_score == pytest.approx(1.0 + FUSION_BONUS * 0.95) + assert alpha.dual_source is True + + repository._rerank_provider = _ReverseReranker() + repository._semantic_vector_k = 3 + repository._reranker_candidates = 2 + reranked_trace = await explain_query( + search_service, + SearchQuery(text="auth", retrieval_mode=SearchRetrievalMode.VECTOR), + limit=3, + offset=0, + ) + assert isinstance(reranked_trace, VectorQueryTrace) + assert reranked_trace.rerank is not None + assert reranked_trace.rerank.stable_pool_refetched is True + assert reranked_trace.rerank.entries[0].key == ("entity", 3) + alpha_rerank = next( + entry for entry in reranked_trace.rerank.entries if entry.key == ("entity", 1) + ) + assert alpha_rerank.pre_rerank_rank == 1 + assert alpha_rerank.pre_rerank_score == pytest.approx(0.95) + assert alpha_rerank.post_rerank_rank == 2 + tail = reranked_trace.rerank.entries[-1] + assert tail.rerank_score is None + assert tail.demoted_score is not None + assert tail.demoted_score < reranked_trace.rerank.tail_floor + + +@pytest.mark.asyncio +async def test_external_overfetch_trace_trims_to_the_candidate_window( + session_maker, + test_project, + app_config, +): + """Chunks hydrated beyond candidate_limit never ran; the trace must not invent them.""" + repository, vector_index = _repository(session_maker, test_project, app_config) + await repository.init_search_index() + live_ids = list(range(1, 21)) + await repository.bulk_index_items( + [ + _row(repository.project_id, entity_id, f"Note {entity_id}", "keep") + for entity_id in live_ids + ] + ) + configured_model = repository._embedding_model_key() + async with db.scoped_session(repository.session_maker) as session: + await session.execute( + text( + "INSERT INTO search_vector_chunks (" + "project_id, entity_id, chunk_key, chunk_text, source_hash, " + "entity_fingerprint, embedding_model, vector_index, embedding_status" + ") VALUES (" + ":project_id, :entity_id, :chunk_key, :chunk_text, 'hash', 'fingerprint', " + ":embedding_model, 'trace-test', 'ready')" + ), + [ + { + "project_id": repository.project_id, + "entity_id": entity_id, + "chunk_key": f"entity:{entity_id}:0", + "chunk_text": f"Note {entity_id} auth retrieval", + "embedding_model": configured_model, + } + for entity_id in live_ids + ], + ) + await session.commit() + + # A second owner sharing entity 1's parseable chunk key — permitted storage shape + # since manifest uniqueness includes entity_id. It hydrates beyond the window and + # must be trimmed even though the key string matches an in-window chunk. + async with db.scoped_session(repository.session_maker) as session: + await session.execute( + text( + "INSERT INTO search_vector_chunks (" + "project_id, entity_id, chunk_key, chunk_text, source_hash, " + "entity_fingerprint, embedding_model, vector_index, embedding_status" + ") VALUES (" + ":project_id, 999, 'entity:1:0', 'Stale twin', 'hash', 'fingerprint', " + ":embedding_model, 'trace-test', 'ready')" + ), + {"project_id": repository.project_id, "embedding_model": configured_model}, + ) + await session.commit() + + # First scan (limit 10): five stale hits crowd out live rows, forcing the + # geometric rescan; the second scan hydrates 14 live chunks plus the foreign-owner + # twin for a 10-wide window. + stale = [ + VectorMatch(VectorKey(100 + offset, f"entity:{100 + offset}:0"), 0.99 - offset * 0.001) + for offset in range(5) + ] + live = [ + VectorMatch(VectorKey(entity_id, f"entity:{entity_id}:0"), 0.9 - entity_id * 0.001) + for entity_id in live_ids + ] + vector_index.matches = stale + live[:14] + [VectorMatch(VectorKey(999, "entity:1:0"), 0.5)] + + collector = SearchTraceCollector() + results = await repository.search( + search_text="auth", + retrieval_mode=SearchRetrievalMode.VECTOR, + min_similarity=0.0, + limit=1, + trace=collector, + ) + + assert len(results) == 1 + assert collector.vector is not None + assert collector.vector.adapter_match_count == 20 + # Hydration accounting keeps full-scan scope (15 ready chunks, 5 stale drops); + # only the served chunk matches are trimmed to the 10-wide candidate window. + assert collector.vector.hydrated_count == 15 + assert len(collector.vector.drops) == 5 + served = {chunk_match.key for chunk_match in collector.vector.chunk_matches} + assert served == {("entity", entity_id) for entity_id in live_ids[:10]} + assert {drop.reason for drop in collector.vector.drops} == {"not_in_manifest"} + + # The foreign-owner twin hydrated beyond the window and must not survive the trim + # under entity 1's identity. + assert all(chunk_match.entity_id != 999 for chunk_match in collector.vector.chunk_matches) + + trace = finalize_query_trace( + collector, + replace(_meta("vector"), limit=1), + ( + FinalResultEntry( + ("entity", live_ids[0]), + live_ids[0], + f"Note {live_ids[0]}", + f"notes/{live_ids[0]}", + f"{live_ids[0]}.md", + 1, + 0.9, + ), + ), + "vector", + ) + response = query_trace_response(trace) + # The truncation is a stage of its own: vector 20→15 (5 hydration drops), then + # candidate_window 15→10, so no count ever mixes scopes. + vector_stage = next(stage for stage in response.stages if stage.name == "vector") + assert (vector_stage.count_in, vector_stage.count_out, vector_stage.dropped) == (20, 15, 5) + window = next(stage for stage in response.stages if stage.name == "candidate_window") + assert (window.count_in, window.count_out, window.dropped) == (15, 10, 5) + row_collapse = next(stage for stage in response.stages if stage.name == "row_collapse") + assert row_collapse.count_in == 10 + + +@pytest.mark.asyncio +async def test_hybrid_trace_drops_vector_rows_cut_before_fusion( + session_maker, + test_project, + app_config, +): + """Rows the vector leg hydrates but never hands to fusion must not read as returned.""" + repository, vector_index = _repository(session_maker, test_project, app_config) + await repository.init_search_index() + live_ids = list(range(1, 21)) + await repository.bulk_index_items( + [ + _row(repository.project_id, entity_id, f"Note {entity_id}", "keep") + for entity_id in live_ids + ] + ) + configured_model = repository._embedding_model_key() + async with db.scoped_session(repository.session_maker) as session: + await session.execute( + text( + "INSERT INTO search_vector_chunks (" + "project_id, entity_id, chunk_key, chunk_text, source_hash, " + "entity_fingerprint, embedding_model, vector_index, embedding_status" + ") VALUES (" + ":project_id, :entity_id, :chunk_key, :chunk_text, 'hash', 'fingerprint', " + ":embedding_model, 'trace-test', 'ready')" + ), + [ + { + "project_id": repository.project_id, + "entity_id": entity_id, + "chunk_key": f"entity:{entity_id}:0", + "chunk_text": f"Note {entity_id} auth retrieval", + "embedding_model": configured_model, + } + for entity_id in live_ids + ], + ) + await session.commit() + vector_index.matches = [ + VectorMatch(VectorKey(entity_id, f"entity:{entity_id}:0"), 0.9 - entity_id * 0.001) + for entity_id in live_ids + ] + + collector = SearchTraceCollector() + results = await repository.search( + search_text="auth", + retrieval_mode=SearchRetrievalMode.HYBRID, + min_similarity=0.0, + limit=1, + trace=collector, + ) + trace = finalize_query_trace( + collector, + replace(_meta("hybrid"), limit=1), + tuple( + FinalResultEntry( + (row.type, row.id), + row.entity_id, + row.title, + row.permalink, + row.file_path, + rank, + row.score or 0.0, + ) + for rank, row in enumerate(results, start=1) + ), + "hybrid", + ) + + response = query_trace_response(trace) + + # Every "returned" candidate must correspond to an actual final result; rows the + # vector leg hydrated but cut before fusion may not masquerade as returned. + returned = [c for c in response.candidates if c.disposition == "returned"] + assert len(returned) == len(results) == 1 + assert collector.vector is not None + traced_vector_rows = {chunk_match.key for chunk_match in collector.vector.chunk_matches} + assert collector.fusion is not None + fused_keys = {entry.key for entry in collector.fusion.entries} + assert traced_vector_rows <= fused_keys + + +@pytest.mark.asyncio +async def test_hydrated_chunk_with_malformed_key_is_an_explicit_drop( + session_maker, + test_project, + app_config, + engine_factory, +): + """A ready chunk with an unparseable key must be named in the trace, not vanish.""" + repository, vector_index = _repository(session_maker, test_project, app_config) + await _seed_trace_corpus(repository, vector_index) + configured_model = repository._embedding_model_key() + async with db.scoped_session(repository.session_maker) as session: + await session.execute( + text( + "INSERT INTO search_vector_chunks (" + "project_id, entity_id, chunk_key, chunk_text, source_hash, " + "entity_fingerprint, embedding_model, vector_index, embedding_status" + ") VALUES (" + ":project_id, 1, 'orphan-key-without-shape', 'Orphan text', 'hash', " + "'fingerprint', :embedding_model, 'trace-test', 'ready')" + ), + {"project_id": repository.project_id, "embedding_model": configured_model}, + ) + vector_index.matches = [ + VectorMatch(VectorKey(1, "entity:1:0"), 0.95), + VectorMatch(VectorKey(1, "orphan-key-without-shape"), 0.9), + ] + + collector = SearchTraceCollector() + results = await repository.search( + search_text="auth", + note_types=["keep"], + retrieval_mode=SearchRetrievalMode.VECTOR, + min_similarity=0.5, + limit=10, + trace=collector, + ) + + assert [(row.type, row.id) for row in results] == [("entity", 1)] + assert collector.vector is not None + malformed = [drop for drop in collector.vector.drops if drop.reason == "malformed_key"] + assert [(drop.chunk_key, drop.similarity) for drop in malformed] == [ + ("orphan-key-without-shape", 0.9) + ] + assert all( + chunk_match.chunk_key != "orphan-key-without-shape" + for chunk_match in collector.vector.chunk_matches + ) + # Malformed hits are drops, not output: in=2 adapter matches, out=1 served chunk. + assert collector.vector.adapter_match_count == 2 + assert collector.vector.hydrated_count == 1 + + +@pytest.mark.asyncio +async def test_explain_query_without_criteria_fails_fast( + session_maker, + test_project, + app_config, +): + """An unsatisfiable query never executes, so explain refuses to invent a trace.""" + repository, _vector_index = _repository(session_maker, test_project, app_config) + search_service = SearchService(repository, MagicMock(), MagicMock(), session_maker) + + with pytest.raises(ValueError, match="without an FTS stage"): + await explain_query(search_service, SearchQuery(), limit=10, offset=0) + + +@pytest.mark.asyncio +async def test_empty_page_distinguishes_pagination_from_empty_candidates( + session_maker, + test_project, + app_config, +): + """An offset past the ranked rows is page_out_of_range, not no_candidates.""" + repository, vector_index = _repository(session_maker, test_project, app_config) + await _seed_trace_corpus(repository, vector_index) + repository._rerank_provider = _ReverseReranker() + search_service = SearchService(repository, MagicMock(), MagicMock(), session_maker) + + beyond_page = await explain_query( + search_service, + SearchQuery(text="auth", retrieval_mode=SearchRetrievalMode.VECTOR), + limit=3, + offset=50, + ) + assert beyond_page.final == () + assert beyond_page.meta.rerank_applied is False + assert beyond_page.meta.rerank_skipped_reason == "page_out_of_range" + + # Hydrated chunks whose rows all fall to the threshold are an empty ranked set: + # offset 0 cannot be "out of range", whatever the hydration count says. + all_rejected = await explain_query( + search_service, + SearchQuery( + text="auth", + retrieval_mode=SearchRetrievalMode.VECTOR, + min_similarity=0.99, + ), + limit=3, + offset=0, + ) + assert all_rejected.final == () + assert all_rejected.meta.rerank_skipped_reason == "no_candidates" + + vector_index.matches = [] + empty_retrieval = await explain_query( + search_service, + SearchQuery(text="auth", retrieval_mode=SearchRetrievalMode.VECTOR), + limit=3, + offset=0, + ) + assert empty_retrieval.final == () + assert empty_retrieval.meta.rerank_skipped_reason == "no_candidates" + + +def test_negative_similarities_survive_without_zero_clamping(): + """With no threshold, negatively correlated matches keep their real scores.""" + vector = build_vector_stage( + candidate_limit=10, + adapter_match_count=3, + hydrated_count=2, + drops=(HydrationDropped(11, "malformed-key", -0.2, "not_in_manifest", None, None),), + chunk_matches={ + ("entity", 1): [("entity:1:0", -0.5, 1), ("entity:1:1", -0.4, 1)], + }, + effective_min_similarity=0.0, + ) + collector = SearchTraceCollector( + vector=vector, + readiness=ManifestReadiness("i", "m", 2, 0, 0), + ) + final = (FinalResultEntry(("entity", 1), 1, "One", "one", "one.md", 1, -0.4),) + trace = finalize_query_trace(collector, _meta("vector"), final, "vector") + + response = query_trace_response(trace) + + by_id = {candidate.id: candidate for candidate in response.candidates} + assert by_id[1].scores.vector_similarity == pytest.approx(-0.4) + malformed = next(candidate for candidate in response.candidates if candidate.id is None) + assert malformed.scores.vector_similarity == pytest.approx(-0.2) + + +def test_relaxed_fts_fallback_is_serialized_on_the_stage(): + """Results produced by the relaxed retry must not read as strict-query matches.""" + relaxed = build_fts_page_stage( + [(("entity", 1), -1.0)], + relaxed_fallback_used=True, + ) + trace = finalize_query_trace( + SearchTraceCollector(fts=relaxed), + _meta("fts"), + (FinalResultEntry(("entity", 1), 1, "One", "one", "one.md", 1, -1.0),), + "fts", + ) + + response = query_trace_response(trace) + + fts_stage = next(stage for stage in response.stages if stage.name == "fts") + assert fts_stage.relaxed_fallback_used is True + + +def test_parseable_drop_with_foreign_owner_stays_a_distinct_candidate(): + """A stale hit under a reused row id must not merge with the live row's evidence.""" + vector = build_vector_stage( + candidate_limit=10, + adapter_match_count=3, + hydrated_count=1, + drops=( + # Live owner's sibling chunk drop merges with the live candidate as before. + HydrationDropped(5, "entity:5:1", 0.6, "pending", "m", "i"), + # Foreign owner: same parseable chunk key, different entity — a stale + # adapter hit surviving a search-row ID reuse. + HydrationDropped(99, "entity:5:0", 0.55, "not_in_manifest", None, None), + ), + chunk_matches={("entity", 5): [("entity:5:0", 0.9, 5)]}, + ) + collector = SearchTraceCollector( + vector=vector, + readiness=ManifestReadiness("i", "m", 1, 1, 0), + ) + final = (FinalResultEntry(("entity", 5), 5, "Five", "five", "five.md", 1, 0.9),) + trace = finalize_query_trace(collector, _meta("vector"), final, "vector") + + response = query_trace_response(trace, {5: "external-5", 99: "external-99"}) + + returned = [c for c in response.candidates if c.disposition == "returned"] + assert [(c.id, c.external_id) for c in returned] == [(5, "external-5")] + live = returned[0] + # The live row keeps only its own evidence: one hydrated chunk + its own sibling drop. + assert [chunk.chunk_key for chunk in live.matched_chunks] == ["entity:5:0"] + assert [chunk.chunk_key for chunk in live.dropped_chunks] == ["entity:5:1"] + foreign = next(c for c in response.candidates if c.external_id == "external-99") + assert foreign.disposition == "not_in_manifest" + assert foreign.scores.vector_similarity == pytest.approx(0.55) + assert [chunk.chunk_key for chunk in foreign.dropped_chunks] == ["entity:5:0"] + + +def test_foreign_owner_drop_stays_split_from_fts_returned_row(): + """A stale same-key drop must not attach to a row returned via FTS alone.""" + fts = build_fts_page_stage( + [(("entity", 5), -1.0)], + normalized_scores={("entity", 5): 1.0}, + fts_max_abs=1.0, + relaxed_fallback_used=False, + ) + vector = build_vector_stage( + candidate_limit=10, + adapter_match_count=1, + hydrated_count=0, + drops=(HydrationDropped(99, "entity:5:0", 0.5, "not_in_manifest", None, None),), + ) + fusion = build_fusion_stage( + formula_version="max+0.3*min/v1", + bonus=FUSION_BONUS, + fts_scores={("entity", 5): 1.0}, + fts_ranks={("entity", 5): 0}, + vector_scores={}, + vector_ranks={}, + ranked_scores=[(("entity", 5), 1.0)], + fusion_ms=0.1, + ) + trace = finalize_query_trace( + SearchTraceCollector( + fts=fts, + vector=vector, + fusion=fusion, + readiness=ManifestReadiness("i", "m", 0, 0, 1), + ), + _meta("hybrid"), + (FinalResultEntry(("entity", 5), 5, "Five", "five", "five.md", 1, 1.0),), + "hybrid", + ) + + response = query_trace_response(trace, {5: "external-5", 99: "external-99"}) + + returned = [c for c in response.candidates if c.disposition == "returned"] + assert [(c.id, c.external_id) for c in returned] == [(5, "external-5")] + assert returned[0].dropped_chunks == [] + foreign = next(c for c in response.candidates if c.external_id == "external-99") + assert foreign.disposition == "not_in_manifest" + assert [chunk.chunk_key for chunk in foreign.dropped_chunks] == ["entity:5:0"] + + +def test_identically_malformed_keys_from_two_entities_stay_distinct_candidates(): + """Two owners serving the same malformed key must not merge their drop evidence.""" + vector = build_vector_stage( + candidate_limit=10, + adapter_match_count=3, + hydrated_count=1, + drops=( + HydrationDropped(11, "malformed-key", 0.35, "not_in_manifest", None, None), + HydrationDropped(12, "malformed-key", 0.30, "not_in_manifest", None, None), + ), + chunk_matches={("entity", 1): [("entity:1:0", 0.9, 1)]}, + ) + collector = SearchTraceCollector( + vector=vector, + readiness=ManifestReadiness("i", "m", 1, 0, 0), + ) + final = (FinalResultEntry(("entity", 1), 1, "One", "one", "one.md", 1, 0.9),) + trace = finalize_query_trace(collector, _meta("vector"), final, "vector") + + response = query_trace_response(trace, {11: "external-11", 12: "external-12"}) + + malformed = sorted( + (c for c in response.candidates if c.id is None), + key=lambda candidate: candidate.external_id or "", + ) + assert [ + (c.external_id, [chunk.chunk_key for chunk in c.dropped_chunks]) for c in malformed + ] == [ + ("external-11", ["malformed-key"]), + ("external-12", ["malformed-key"]), + ] + assert {c.external_id: c.scores.vector_similarity for c in malformed} == { + "external-11": pytest.approx(0.35), + "external-12": pytest.approx(0.30), + } + + +def test_duplicate_final_rows_keep_one_returned_candidate_per_occurrence(): + """SQLite FTS duplicates keep per-occurrence stage ranks and final ranks.""" + duplicated_page = build_fts_page_stage( + [(("entity", 1), -1.0), (("entity", 1), -1.0), (("entity", 2), -0.5)], + relaxed_fallback_used=False, + ) + # The stage freezes the executed page occurrence-by-occurrence: collapsing it + # would undercount the page and shift the later unique row's source rank. + assert duplicated_page.result_count == 3 + assert [(score.key, score.rank) for score in duplicated_page.raw_scores] == [ + (("entity", 1), 1), + (("entity", 1), 2), + (("entity", 2), 3), + ] + + collector = SearchTraceCollector(fts=duplicated_page) + final = ( + FinalResultEntry(("entity", 1), 1, "One", "one", "one.md", 1, 0.9), + FinalResultEntry(("entity", 1), 1, "One", "one", "one.md", 2, 0.9), + FinalResultEntry(("entity", 2), 2, "Two", "two", "two.md", 3, 0.4), + ) + trace = finalize_query_trace(collector, _meta("fts"), final, "fts") + + response = query_trace_response(trace, {1: "external-1", 2: "external-2"}) + + returned = [c for c in response.candidates if c.disposition == "returned"] + assert [(c.id, c.scores.final_rank) for c in returned] == [(1, 1), (1, 2), (2, 3)] + # The logical candidate keeps its best (first-occurrence) source rank, and the + # later unique row keeps the rank of its actual page position. + assert [c.scores.fts_rank for c in returned] == [1, 1, 3] + + +def test_non_text_criteria_and_null_owner_rows_stay_inspectable(): + """Prepared criteria render their executed form; NULL-owner rows degrade, not reject.""" + from basic_memory.services.search_service import PreparedSearchQuery, describe_search_criteria + + empty = PreparedSearchQuery( + search_text=None, + permalink=None, + permalink_match=None, + title=None, + note_types=None, + search_item_types=None, + categories=None, + after_date=None, + metadata_filters=None, + retrieval_mode=SearchRetrievalMode.FTS, + min_similarity=None, + ) + + described = describe_search_criteria(replace(empty, title="Search Spec", note_types=["note"])) + assert described == "title=\"Search Spec\" note_types=['note']" + + # SearchItemType's str() is "SearchItemType.ENTITY"; the executed filter is "entity". + typed = describe_search_criteria( + replace(empty, search_item_types=[SearchItemType.ENTITY, SearchItemType.OBSERVATION]) + ) + assert typed == "entity_types=['entity', 'observation']" + + # tags=/status= shorthand folds into metadata_filters during prepare_query, so the + # description shows the executed filters rather than the raw request fields. + filter_only = describe_search_criteria( + replace( + empty, + categories=["decision"], + metadata_filters={"tags": ["auth"], "status": "active", "priority": "high"}, + ) + ) + assert filter_only == ( + "categories=['decision'] " + "metadata_filters={'tags': ['auth'], 'status': 'active', 'priority': 'high'}" + ) + + degenerate = build_fts_page_stage([(("entity", 7), -1.0)], relaxed_fallback_used=False) + collector = SearchTraceCollector(fts=degenerate) + final = (FinalResultEntry(("entity", 7), None, "Legacy", None, "legacy.md", 1, 0.5),) + trace = finalize_query_trace(collector, _meta("fts"), final, "fts") + + response = query_trace_response(trace, {}) + + returned = [c for c in response.candidates if c.disposition == "returned"] + assert [(c.id, c.external_id) for c in returned] == [(7, None)] diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 70d44e425..57a1a8d9b 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -18,6 +18,7 @@ SearchRepositoryBase, _PreparedEntityVectorSync, ) +from basic_memory.repository.search_trace import SearchTraceCollector from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository from basic_memory.repository.semantic_errors import ( SemanticSearchDisabledError, @@ -91,6 +92,8 @@ async def search( limit: int = 10, offset: int = 0, allow_relaxed: bool = False, + *, + trace: SearchTraceCollector | None = None, ) -> list[SearchIndexRow]: return [] @@ -99,7 +102,14 @@ async def _ensure_vector_tables(self): pass @override - async def _run_vector_query(self, session, query_embedding, candidate_limit): + async def _run_vector_query( + self, + session, + query_embedding, + candidate_limit, + *, + trace: SearchTraceCollector | None = None, + ): return [] @override diff --git a/tests/repository/test_semantic_vector_sync.py b/tests/repository/test_semantic_vector_sync.py index 465ded2a6..0e6a8c12e 100644 --- a/tests/repository/test_semantic_vector_sync.py +++ b/tests/repository/test_semantic_vector_sync.py @@ -13,6 +13,7 @@ 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_repository_base import SearchRepositoryBase +from basic_memory.repository.search_trace import SearchTraceCollector from basic_memory.repository.semantic_chunking import VectorChunkRecord from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode @@ -60,6 +61,8 @@ async def search( limit: int = 10, offset: int = 0, allow_relaxed: bool = False, + *, + trace: SearchTraceCollector | None = None, ) -> list[SearchIndexRow]: return [] @@ -68,7 +71,14 @@ async def _ensure_vector_tables(self): pass @override - async def _run_vector_query(self, session, query_embedding, candidate_limit): + async def _run_vector_query( + self, + session, + query_embedding, + candidate_limit, + *, + trace: SearchTraceCollector | None = None, + ): return [] @override diff --git a/tests/repository/test_vector_pagination.py b/tests/repository/test_vector_pagination.py index 0764c404b..004969de1 100644 --- a/tests/repository/test_vector_pagination.py +++ b/tests/repository/test_vector_pagination.py @@ -14,6 +14,7 @@ from basic_memory.repository.search_repository_base import 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 @@ -70,6 +71,8 @@ async def search( limit: int = 10, offset: int = 0, allow_relaxed: bool = False, + *, + trace: SearchTraceCollector | None = None, ) -> list[SearchIndexRow]: return [] # pragma: no cover @@ -78,7 +81,14 @@ async def _ensure_vector_tables(self): pass # pragma: no cover @override - async def _run_vector_query(self, session, query_embedding, candidate_limit): + async def _run_vector_query( + self, + session, + query_embedding, + candidate_limit, + *, + trace: SearchTraceCollector | None = None, + ): return [] # pragma: no cover @override diff --git a/tests/repository/test_vector_threshold.py b/tests/repository/test_vector_threshold.py index dcd5db25c..cbe55ad54 100644 --- a/tests/repository/test_vector_threshold.py +++ b/tests/repository/test_vector_threshold.py @@ -15,6 +15,7 @@ TOP_CHUNKS_PER_RESULT, SearchRepositoryBase, ) +from basic_memory.repository.search_trace import SearchTraceCollector from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode @@ -74,6 +75,8 @@ async def search( limit: int = 10, offset: int = 0, allow_relaxed: bool = False, + *, + trace: SearchTraceCollector | None = None, ) -> list[SearchIndexRow]: return [] # pragma: no cover @@ -82,7 +85,14 @@ async def _ensure_vector_tables(self): pass # pragma: no cover @override - async def _run_vector_query(self, session, query_embedding, candidate_limit): + async def _run_vector_query( + self, + session, + query_embedding, + candidate_limit, + *, + trace: SearchTraceCollector | None = None, + ): return [] # pragma: no cover @override diff --git a/tests/services/test_note_type_normalization.py b/tests/services/test_note_type_normalization.py index c8ba3d725..834befc10 100644 --- a/tests/services/test_note_type_normalization.py +++ b/tests/services/test_note_type_normalization.py @@ -13,6 +13,7 @@ from basic_memory.repository.search_index_row import SearchIndexRow from basic_memory.repository.search_repository import SearchRepository from basic_memory.schemas.search import SearchQuery +from basic_memory.services.retrieval_inspect import explain_query from basic_memory.services.search_service import SearchService @@ -25,13 +26,13 @@ def _search_service(repository: SearchRepository) -> SearchService: ) -def test_prepare_query_canonicalizes_directly_assigned_note_types(): +def testprepare_query_canonicalizes_directly_assigned_note_types(): """Service callers cannot bypass canonicalization by mutating SearchQuery.""" repository = cast(SearchRepository, MagicMock()) service = _search_service(repository) query = SearchQuery.model_construct(note_types=["TaskItem"]) - prepared = service._prepare_query(query) + prepared = service.prepare_query(query) assert prepared is not None assert prepared.note_types == ["task_item"] @@ -112,3 +113,10 @@ async def test_search_matches_legacy_note_type_projection_without_reindex( assert [result.entity_id for result in results] == [entity.id] assert await search_service.count(query) == 1 + + # The execution trace must describe the expanded filter that admitted the legacy + # row, not the raw request — a trace showing only ['task_item'] could not explain + # why a 'TaskItem' row was returned. + trace = await explain_query(search_service, query, limit=10, offset=0) + assert "note_types=['TaskItem', 'task_item']" in trace.meta.query_text + assert [entry.entity_id for entry in trace.final] == [entity.id] diff --git a/tests/services/test_search_service.py b/tests/services/test_search_service.py index 75df0ae49..5398efc7a 100644 --- a/tests/services/test_search_service.py +++ b/tests/services/test_search_service.py @@ -295,7 +295,7 @@ async def test_search_entity_type(search_service, test_graph): @pytest.mark.asyncio async def test_search_categories_filter(search_service, test_graph): - """categories propagates through _prepare_query/has_criteria to scope results. + """categories propagates through prepare_query/has_criteria to scope results. The test_graph fixture indexes observations with categories "note" and "tech". A categories filter must return only matching observation categories.