diff --git a/src/basic_memory/api/v2/routers/search_router.py b/src/basic_memory/api/v2/routers/search_router.py index 4c8f2f9b..be80d4b4 100644 --- a/src/basic_memory/api/v2/routers/search_router.py +++ b/src/basic_memory/api/v2/routers/search_router.py @@ -4,6 +4,8 @@ This router uses external_id UUIDs for stable, API-friendly routing. V1 uses string-based project names which are less efficient and less stable. """ +import asyncio + from fastapi import APIRouter, HTTPException, Path import logfire @@ -12,7 +14,7 @@ from basic_memory.repository.semantic_errors import ( SemanticDependenciesMissingError, SemanticSearchDisabledError, ) -from basic_memory.schemas.search import SearchQuery, SearchResponse +from basic_memory.schemas.search import SearchQuery, SearchResponse, SearchRetrievalMode from basic_memory.deps import ( SearchServiceV2ExternalDep, EntityServiceV2ExternalDep, @@ -65,7 +67,7 @@ async def search( has_filters=bool(query.note_types or query.entity_types or query.metadata_filters), ): offset = (page - 1) * page_size - fetch_limit = page_size + 1 + exact_count_available = query.retrieval_mode == SearchRetrievalMode.FTS try: with logfire.span( "api.search.search.execute_query", @@ -75,7 +77,14 @@ async def search( page=page, page_size=page_size, ): - results = await search_service.search(query, limit=fetch_limit, offset=offset) + if exact_count_available: + results, total = await asyncio.gather( + search_service.search(query, limit=page_size, offset=offset), + search_service.count(query), + ) + else: + results = await search_service.search(query, limit=page_size + 1, offset=offset) + total = 0 except SemanticSearchDisabledError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc except SemanticDependenciesMissingError as exc: @@ -90,9 +99,15 @@ async def search( phase="paginate_results", result_count=len(results), ): - has_more = len(results) > page_size - if has_more: - results = results[:page_size] + if exact_count_available: + has_more = offset + len(results) < total + else: + # Trigger: semantic modes would need another vector/hybrid retrieval to count. + # Why: search requests should not pay for a second semantic pass. + # Outcome: preserve probe pagination for semantic search and leave total at 0. + has_more = len(results) > page_size + if has_more: + results = results[:page_size] with logfire.span( "api.search.search.hydrate_results", @@ -113,6 +128,7 @@ async def search( results=search_results, current_page=page, page_size=page_size, + total=total, has_more=has_more, ) diff --git a/src/basic_memory/repository/postgres_search_repository.py b/src/basic_memory/repository/postgres_search_repository.py index eb2ffa86..8e44f371 100644 --- a/src/basic_memory/repository/postgres_search_repository.py +++ b/src/basic_memory/repository/postgres_search_repository.py @@ -686,7 +686,17 @@ class PostgresSearchRepository(SearchRepositoryBase): # FTS search (Postgres-specific) # ------------------------------------------------------------------ - async def search( + @staticmethod + def _is_tsquery_syntax_error(exc: Exception) -> bool: + msg = str(exc).lower() + return ( + "syntax error in tsquery" in msg + or "invalid input syntax for type tsquery" in msg + or "no operand in tsquery" in msg + or "no operator in tsquery" in msg + ) + + async def _build_fts_query_parts( self, search_text: Optional[str] = None, permalink: Optional[str] = None, @@ -696,31 +706,8 @@ class PostgresSearchRepository(SearchRepositoryBase): after_date: Optional[datetime] = None, search_item_types: Optional[List[SearchItemType]] = None, metadata_filters: Optional[dict] = None, - retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, - min_similarity: Optional[float] = None, - limit: int = 10, - offset: int = 0, - ) -> List[SearchIndexRow]: - """Search across all indexed content using PostgreSQL tsvector.""" - # --- Dispatch vector / hybrid modes (shared logic) --- - dispatched = await self._dispatch_retrieval_mode( - search_text=search_text, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - metadata_filters=metadata_filters, - retrieval_mode=retrieval_mode, - min_similarity=min_similarity, - limit=limit, - offset=offset, - ) - if dispatched is not None: - return dispatched - - # --- FTS mode (Postgres-specific) --- + ) -> tuple[str, str, dict, str, str]: + """Build Postgres FTS FROM/WHERE params shared by search and count.""" conditions = [] params = {} order_by_clause = "" @@ -868,10 +855,6 @@ class PostgresSearchRepository(SearchRepositoryBase): params["project_id"] = self.project_id conditions.append("search_index.project_id = :project_id") - # set limit and offset - params["limit"] = limit - params["offset"] = offset - # Build WHERE clause where_clause = " AND ".join(conditions) if conditions else "1=1" @@ -884,6 +867,64 @@ class PostgresSearchRepository(SearchRepositoryBase): else: score_expr = "0" + return from_clause, where_clause, params, order_by_clause, score_expr + + async def search( + self, + search_text: Optional[str] = None, + permalink: Optional[str] = None, + permalink_match: Optional[str] = None, + title: Optional[str] = None, + note_types: Optional[List[str]] = None, + after_date: Optional[datetime] = None, + search_item_types: Optional[List[SearchItemType]] = None, + metadata_filters: Optional[dict] = None, + retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, + min_similarity: Optional[float] = None, + limit: int = 10, + offset: int = 0, + ) -> List[SearchIndexRow]: + """Search across all indexed content using PostgreSQL tsvector.""" + # --- Dispatch vector / hybrid modes (shared logic) --- + dispatched = await self._dispatch_retrieval_mode( + search_text=search_text, + permalink=permalink, + permalink_match=permalink_match, + title=title, + note_types=note_types, + after_date=after_date, + search_item_types=search_item_types, + metadata_filters=metadata_filters, + retrieval_mode=retrieval_mode, + min_similarity=min_similarity, + limit=limit, + offset=offset, + ) + if dispatched is not None: + return dispatched + + # --- FTS mode (Postgres-specific) --- + ( + from_clause, + where_clause, + params, + order_by_clause, + score_expr, + ) = await self._build_fts_query_parts( + search_text=search_text, + permalink=permalink, + permalink_match=permalink_match, + title=title, + note_types=note_types, + after_date=after_date, + search_item_types=search_item_types, + metadata_filters=metadata_filters, + ) + + # set limit and offset + params["limit"] = limit + params["offset"] = offset + sql = f""" SELECT search_index.project_id, @@ -915,17 +956,7 @@ class PostgresSearchRepository(SearchRepositoryBase): result = await session.execute(text(sql), params) rows = result.fetchall() except Exception as e: - # Handle tsquery syntax errors (and only those). - # - # Important: Postgres errors for other failures (e.g. missing table) will still mention - # `to_tsquery(...)` in the SQL text, so checking for the substring "tsquery" is too broad. - msg = str(e).lower() - if ( - "syntax error in tsquery" in msg - or "invalid input syntax for type tsquery" in msg - or "no operand in tsquery" in msg - or "no operator in tsquery" in msg - ): + if self._is_tsquery_syntax_error(e): logger.warning(f"tsquery syntax error for search term: {search_text}, error: {e}") return [] @@ -966,3 +997,60 @@ class PostgresSearchRepository(SearchRepositoryBase): ) return results + + async def count( + self, + search_text: Optional[str] = None, + permalink: Optional[str] = None, + permalink_match: Optional[str] = None, + title: Optional[str] = None, + note_types: Optional[List[str]] = None, + after_date: Optional[datetime] = None, + search_item_types: Optional[List[SearchItemType]] = None, + metadata_filters: Optional[dict] = None, + retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, + min_similarity: Optional[float] = None, + ) -> int: + """Count indexed content matching the Postgres FTS query.""" + if retrieval_mode != SearchRetrievalMode.FTS: + return await super().count( + search_text=search_text, + permalink=permalink, + permalink_match=permalink_match, + title=title, + note_types=note_types, + after_date=after_date, + search_item_types=search_item_types, + metadata_filters=metadata_filters, + retrieval_mode=retrieval_mode, + min_similarity=min_similarity, + ) + + ( + from_clause, + where_clause, + params, + _order_by_clause, + _score_expr, + ) = await self._build_fts_query_parts( + search_text=search_text, + permalink=permalink, + permalink_match=permalink_match, + title=title, + note_types=note_types, + after_date=after_date, + search_item_types=search_item_types, + metadata_filters=metadata_filters, + ) + sql = f"SELECT COUNT(*) FROM {from_clause} WHERE {where_clause}" + logger.trace(f"Count {sql} params: {params}") + try: + async with db.scoped_session(self.session_maker) as session: + result = await session.execute(text(sql), params) + return int(result.scalar_one()) + except Exception as e: + if self._is_tsquery_syntax_error(e): + logger.warning(f"tsquery syntax error for search term: {search_text}, error: {e}") + return 0 + logger.error(f"Database error during search count: {e}") + raise diff --git a/src/basic_memory/repository/search_repository.py b/src/basic_memory/repository/search_repository.py index d1d2e365..6bfab46e 100644 --- a/src/basic_memory/repository/search_repository.py +++ b/src/basic_memory/repository/search_repository.py @@ -50,6 +50,22 @@ class SearchRepository(Protocol): """Search across indexed content.""" ... + async def count( + self, + search_text: Optional[str] = None, + permalink: Optional[str] = None, + permalink_match: Optional[str] = None, + title: Optional[str] = None, + note_types: Optional[List[str]] = None, + after_date: Optional[datetime] = None, + search_item_types: Optional[List[SearchItemType]] = None, + metadata_filters: Optional[dict] = None, + retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, + min_similarity: Optional[float] = None, + ) -> int: + """Count indexed content matching the same filters as search.""" + ... + async def index_item(self, search_index_row: SearchIndexRow) -> None: """Index a single item.""" ... diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 4459d990..c5367295 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -247,6 +247,24 @@ class SearchRepositoryBase(ABC): """ pass + async def count( + self, + search_text: Optional[str] = None, + permalink: Optional[str] = None, + permalink_match: Optional[str] = None, + title: Optional[str] = None, + note_types: Optional[List[str]] = None, + after_date: Optional[datetime] = None, + search_item_types: Optional[List[SearchItemType]] = None, + metadata_filters: Optional[Dict[str, Any]] = None, + retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, + min_similarity: Optional[float] = None, + ) -> int: + """Count results when a backend-specific COUNT query is available.""" + if retrieval_mode != SearchRetrievalMode.FTS: + raise ValueError("Exact counts are only supported for full-text search retrieval.") + raise NotImplementedError("Backend search repositories must implement full-text counts.") + # ------------------------------------------------------------------ # Abstract methods — semantic search (backend-specific DB operations) # ------------------------------------------------------------------ diff --git a/src/basic_memory/repository/sqlite_search_repository.py b/src/basic_memory/repository/sqlite_search_repository.py index 7d11874e..2eaec06c 100644 --- a/src/basic_memory/repository/sqlite_search_repository.py +++ b/src/basic_memory/repository/sqlite_search_repository.py @@ -701,7 +701,11 @@ class SQLiteSearchRepository(SearchRepositoryBase): # FTS search (backend-specific) # ------------------------------------------------------------------ - async def search( + @staticmethod + def _is_fts5_syntax_error(exc: Exception) -> bool: + return "fts5: syntax error" in str(exc).lower() + + async def _build_fts_query_parts( self, search_text: Optional[str] = None, permalink: Optional[str] = None, @@ -711,31 +715,8 @@ class SQLiteSearchRepository(SearchRepositoryBase): after_date: Optional[datetime] = None, search_item_types: Optional[List[SearchItemType]] = None, metadata_filters: Optional[dict] = None, - retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, - min_similarity: Optional[float] = None, - limit: int = 10, - offset: int = 0, - ) -> List[SearchIndexRow]: - """Search across all indexed content using SQLite FTS5.""" - # --- Dispatch vector / hybrid modes (shared logic) --- - dispatched = await self._dispatch_retrieval_mode( - search_text=search_text, - permalink=permalink, - permalink_match=permalink_match, - title=title, - note_types=note_types, - after_date=after_date, - search_item_types=search_item_types, - metadata_filters=metadata_filters, - retrieval_mode=retrieval_mode, - min_similarity=min_similarity, - limit=limit, - offset=offset, - ) - if dispatched is not None: - return dispatched - - # --- FTS mode (SQLite-specific) --- + ) -> tuple[str, str, dict, str]: + """Build SQLite FTS FROM/WHERE params shared by search and count.""" conditions = [] match_conditions = [] params = {} @@ -911,13 +892,60 @@ class SQLiteSearchRepository(SearchRepositoryBase): params["project_id"] = self.project_id conditions.append("search_index.project_id = :project_id") + # Build WHERE clause + where_clause = " AND ".join(conditions) if conditions else "1=1" + return from_clause, where_clause, params, order_by_clause + + async def search( + self, + search_text: Optional[str] = None, + permalink: Optional[str] = None, + permalink_match: Optional[str] = None, + title: Optional[str] = None, + note_types: Optional[List[str]] = None, + after_date: Optional[datetime] = None, + search_item_types: Optional[List[SearchItemType]] = None, + metadata_filters: Optional[dict] = None, + retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, + min_similarity: Optional[float] = None, + limit: int = 10, + offset: int = 0, + ) -> List[SearchIndexRow]: + """Search across all indexed content using SQLite FTS5.""" + # --- Dispatch vector / hybrid modes (shared logic) --- + dispatched = await self._dispatch_retrieval_mode( + search_text=search_text, + permalink=permalink, + permalink_match=permalink_match, + title=title, + note_types=note_types, + after_date=after_date, + search_item_types=search_item_types, + metadata_filters=metadata_filters, + retrieval_mode=retrieval_mode, + min_similarity=min_similarity, + limit=limit, + offset=offset, + ) + if dispatched is not None: + return dispatched + + # --- FTS mode (SQLite-specific) --- + from_clause, where_clause, params, order_by_clause = await self._build_fts_query_parts( + search_text=search_text, + permalink=permalink, + permalink_match=permalink_match, + title=title, + note_types=note_types, + after_date=after_date, + search_item_types=search_item_types, + metadata_filters=metadata_filters, + ) + # set limit on search query params["limit"] = limit params["offset"] = offset - # Build WHERE clause - where_clause = " AND ".join(conditions) if conditions else "1=1" - sql = f""" SELECT search_index.project_id, @@ -950,7 +978,7 @@ class SQLiteSearchRepository(SearchRepositoryBase): rows = result.fetchall() except Exception as e: # Handle FTS5 syntax errors and provide user-friendly feedback - if "fts5: syntax error" in str(e).lower(): # pragma: no cover + 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 return [] @@ -988,3 +1016,54 @@ class SQLiteSearchRepository(SearchRepositoryBase): ) return results + + async def count( + self, + search_text: Optional[str] = None, + permalink: Optional[str] = None, + permalink_match: Optional[str] = None, + title: Optional[str] = None, + note_types: Optional[List[str]] = None, + after_date: Optional[datetime] = None, + search_item_types: Optional[List[SearchItemType]] = None, + metadata_filters: Optional[dict] = None, + retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, + min_similarity: Optional[float] = None, + ) -> int: + """Count indexed content matching the SQLite FTS query.""" + if retrieval_mode != SearchRetrievalMode.FTS: + return await super().count( + search_text=search_text, + permalink=permalink, + permalink_match=permalink_match, + title=title, + note_types=note_types, + after_date=after_date, + search_item_types=search_item_types, + metadata_filters=metadata_filters, + retrieval_mode=retrieval_mode, + min_similarity=min_similarity, + ) + + from_clause, where_clause, params, _order_by_clause = await self._build_fts_query_parts( + search_text=search_text, + permalink=permalink, + permalink_match=permalink_match, + title=title, + note_types=note_types, + after_date=after_date, + search_item_types=search_item_types, + metadata_filters=metadata_filters, + ) + sql = f"SELECT COUNT(*) FROM {from_clause} WHERE {where_clause}" + logger.trace(f"Count {sql} params: {params}") + try: + async with db.scoped_session(self.session_maker) as session: + result = await session.execute(text(sql), params) + return int(result.scalar_one()) + except Exception as e: + if self._is_fts5_syntax_error(e): # pragma: no cover + logger.warning(f"FTS5 syntax error for search term: {search_text}, error: {e}") + return 0 + logger.error(f"Database error during search count: {e}") + raise diff --git a/src/basic_memory/schemas/search.py b/src/basic_memory/schemas/search.py index a2b94173..51a00d83 100644 --- a/src/basic_memory/schemas/search.py +++ b/src/basic_memory/schemas/search.py @@ -142,4 +142,5 @@ class SearchResponse(BaseModel): results: List[SearchResult] current_page: int page_size: int + total: int = 0 has_more: bool = False diff --git a/src/basic_memory/services/search_service.py b/src/basic_memory/services/search_service.py index 452c7346..46b12ba0 100644 --- a/src/basic_memory/services/search_service.py +++ b/src/basic_memory/services/search_service.py @@ -3,6 +3,7 @@ import asyncio import ast import re +from dataclasses import dataclass from datetime import datetime from typing import List, Optional, Set, Dict, Any @@ -64,6 +65,22 @@ FTS_RELAXED_STOPWORDS = { } +@dataclass(frozen=True) +class _PreparedSearchQuery: + """Normalized query inputs shared by search and count.""" + + search_text: str | None + permalink: str | None + permalink_match: str | None + title: str | None + note_types: list[str] | None + search_item_types: list[SearchItemType] | None + after_date: datetime | None + metadata_filters: dict[str, Any] | None + retrieval_mode: SearchRetrievalMode + min_similarity: float | None + + def _strip_nul(value: str) -> str: """Strip NUL bytes that PostgreSQL text columns cannot store. @@ -132,27 +149,20 @@ class SearchService: logger.info("Reindex complete") - async def search(self, query: SearchQuery, limit=10, offset=0) -> List[SearchIndexRow]: - """Search across all indexed content. + def _prepare_query(self, query: SearchQuery) -> _PreparedSearchQuery | None: + """Normalize a SearchQuery into repository arguments.""" + search_text = query.text + tags = query.tags - Supports three modes: - 1. Exact permalink: finds direct matches for a specific path - 2. Pattern match: handles * wildcards in paths - 3. Text search: full-text search across title/content - """ - # Support tag: shorthand by mapping to tags filter - if query.text: - text = query.text.strip() - if text.lower().startswith("tag:"): - tag_values = re.split(r"[,\s]+", text[4:].strip()) - tags = [t for t in tag_values if t] - if tags: - query.tags = tags - query.text = None - - if query.no_criteria(): - logger.debug("no criteria passed to query") - return [] + # Support tag: shorthand by mapping to tags filter. + if search_text is not None: + search_text = search_text.strip() or None + if search_text and search_text.lower().startswith("tag:"): + tag_values = re.split(r"[,\s]+", search_text[4:].strip()) + parsed_tags = [t for t in tag_values if t] + if parsed_tags: + tags = parsed_tags + search_text = None after_date = ( ( @@ -164,50 +174,124 @@ class SearchService: else None ) - # Merge structured metadata filters (explicit + convenience fields) + # Merge structured metadata filters (explicit + convenience fields). metadata_filters: Optional[Dict[str, Any]] = None - if query.metadata_filters or query.tags or query.status: + if query.metadata_filters or tags or query.status: metadata_filters = dict(query.metadata_filters or {}) - if query.tags: - metadata_filters.setdefault("tags", query.tags) + if tags: + metadata_filters.setdefault("tags", tags) if query.status: metadata_filters.setdefault("status", query.status) - retrieval_mode = query.retrieval_mode or SearchRetrievalMode.FTS - strict_search_text = query.text + prepared = _PreparedSearchQuery( + search_text=search_text, + permalink=query.permalink, + permalink_match=query.permalink_match, + title=query.title, + note_types=query.note_types, + search_item_types=query.entity_types, + after_date=after_date, + metadata_filters=metadata_filters, + retrieval_mode=query.retrieval_mode or SearchRetrievalMode.FTS, + min_similarity=query.min_similarity, + ) + + has_criteria = bool( + prepared.search_text + or prepared.permalink + or prepared.permalink_match + or prepared.title + or prepared.note_types + or prepared.search_item_types + or prepared.after_date + or prepared.metadata_filters + ) + if not has_criteria: + logger.debug("no criteria passed to query") + return None + return prepared + + @staticmethod + def _prepared_has_filters(prepared: _PreparedSearchQuery) -> bool: + return bool( + prepared.metadata_filters + or prepared.note_types + or prepared.search_item_types + or prepared.after_date + ) + + async def _search_repository( + self, + prepared: _PreparedSearchQuery, + *, + search_text: str | None, + limit: int, + offset: int, + ) -> List[SearchIndexRow]: + 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, + after_date=prepared.after_date, + metadata_filters=prepared.metadata_filters, + retrieval_mode=prepared.retrieval_mode, + min_similarity=prepared.min_similarity, + limit=limit, + offset=offset, + ) + + async def _count_repository( + self, + prepared: _PreparedSearchQuery, + *, + search_text: str | None, + ) -> int: + return await self.repository.count( + 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, + after_date=prepared.after_date, + metadata_filters=prepared.metadata_filters, + retrieval_mode=prepared.retrieval_mode, + min_similarity=prepared.min_similarity, + ) + + async def search(self, query: SearchQuery, limit=10, offset=0) -> List[SearchIndexRow]: + """Search across all indexed content. + + Supports three modes: + 1. Exact permalink: finds direct matches for a specific path + 2. Pattern match: handles * wildcards in paths + 3. Text search: full-text search across title/content + """ + prepared = self._prepare_query(query) + if prepared is None: + return [] + + strict_search_text = prepared.search_text has_query = bool( - strict_search_text or query.title or query.permalink or query.permalink_match - ) - has_filters = bool( - metadata_filters - or query.note_types - or query.entity_types - or after_date - or query.tags - or query.status + strict_search_text or prepared.title or prepared.permalink or prepared.permalink_match ) + has_filters = self._prepared_has_filters(prepared) with logfire.span( "search.execute", - retrieval_mode=retrieval_mode.value, + retrieval_mode=prepared.retrieval_mode.value, has_query=has_query, has_filters=has_filters, limit=limit, offset=offset, ): logger.trace(f"Searching with query: {query}") - # First pass: preserve existing strict search behavior. - results = await self.repository.search( + results = await self._search_repository( + prepared, search_text=strict_search_text, - permalink=query.permalink, - permalink_match=query.permalink_match, - title=query.title, - note_types=query.note_types, - search_item_types=query.entity_types, - after_date=after_date, - metadata_filters=metadata_filters, - retrieval_mode=retrieval_mode, - min_similarity=query.min_similarity, limit=limit, offset=offset, ) @@ -217,7 +301,9 @@ class SearchService: # Outcome: retry once with relaxed OR terms while preserving explicit boolean intent. if results: return results - if not self._is_relaxed_fts_fallback_eligible(query, strict_search_text, retrieval_mode): + if not self._is_relaxed_fts_fallback_eligible( + query, strict_search_text, prepared.retrieval_mode + ): return results assert strict_search_text is not None @@ -231,26 +317,57 @@ class SearchService: ) with logfire.span( "search.relaxed_fts_retry", - retrieval_mode=retrieval_mode.value, + retrieval_mode=prepared.retrieval_mode.value, token_count=len(self._tokenize_fts_text(strict_search_text)), limit=limit, offset=offset, ): - return await self.repository.search( + return await self._search_repository( + prepared, search_text=relaxed_search_text, - permalink=query.permalink, - permalink_match=query.permalink_match, - title=query.title, - note_types=query.note_types, - search_item_types=query.entity_types, - after_date=after_date, - metadata_filters=metadata_filters, - retrieval_mode=retrieval_mode, - min_similarity=query.min_similarity, limit=limit, offset=offset, ) + async def count(self, query: SearchQuery) -> int: + """Count all indexed rows matching a query.""" + prepared = self._prepare_query(query) + if prepared is None: + return 0 + + strict_search_text = prepared.search_text + has_query = bool( + strict_search_text or prepared.title or prepared.permalink or prepared.permalink_match + ) + has_filters = self._prepared_has_filters(prepared) + + with logfire.span( + "search.count", + retrieval_mode=prepared.retrieval_mode.value, + has_query=has_query, + has_filters=has_filters, + ): + total = await self._count_repository(prepared, search_text=strict_search_text) + + if total > 0: + return total + if not self._is_relaxed_fts_fallback_eligible( + query, strict_search_text, prepared.retrieval_mode + ): + return total + + assert strict_search_text is not None + relaxed_search_text = self._build_relaxed_fts_query(strict_search_text) + if relaxed_search_text == strict_search_text: + return total + + with logfire.span( + "search.count.relaxed_fts_retry", + retrieval_mode=prepared.retrieval_mode.value, + token_count=len(self._tokenize_fts_text(strict_search_text)), + ): + return await self._count_repository(prepared, search_text=relaxed_search_text) + @staticmethod def _tokenize_fts_text(search_text: str) -> list[str]: """Tokenize text into alphanumeric terms for relaxed FTS fallback.""" diff --git a/tests/api/v2/test_search_router.py b/tests/api/v2/test_search_router.py index 7a8b2bc6..a7f93ce2 100644 --- a/tests/api/v2/test_search_router.py +++ b/tests/api/v2/test_search_router.py @@ -57,7 +57,7 @@ async def test_search_entities( ) # Search for the entity - response = await client.post(f"{v2_project_url}/search/", json={"search_text": "Searchable"}) + response = await client.post(f"{v2_project_url}/search/", json={"text": "Searchable"}) assert response.status_code == 200 data = response.json() @@ -94,7 +94,7 @@ async def test_search_with_pagination( # Search with pagination response = await client.post( f"{v2_project_url}/search/", - json={"search_text": "Search Entity"}, + json={"text": "Search Entity"}, params={"page": 1, "page_size": 3}, ) @@ -102,6 +102,57 @@ async def test_search_with_pagination( data = response.json() assert data["current_page"] == 1 assert data["page_size"] == 3 + assert data["total"] == 5 + assert data["has_more"] is True + + response = await client.post( + f"{v2_project_url}/search/", + json={"text": "Search Entity"}, + params={"page": 2, "page_size": 3}, + ) + + assert response.status_code == 200 + data = response.json() + assert data["current_page"] == 2 + assert data["page_size"] == 3 + assert data["total"] == 5 + assert data["has_more"] is False + assert len(data["results"]) == 2 + + +@pytest.mark.asyncio +async def test_search_with_item_type_filter_returns_total( + client: AsyncClient, + test_project: Project, + v2_project_url: str, + entity_repository, + search_service, + file_service, +): + """Metadata-only graph searches should include exact totals for pagination.""" + for i in range(5): + entity_data = { + "title": f"Structured Entity {i}", + "note_type": "note", + "content_type": "text/markdown", + "file_path": f"structured_{i}.md", + "checksum": f"structuredsum{i}", + } + await create_test_entity( + test_project, entity_data, entity_repository, search_service, file_service + ) + + response = await client.post( + f"{v2_project_url}/search/", + json={"entity_types": ["entity"]}, + params={"page": 1, "page_size": 3}, + ) + + assert response.status_code == 200 + data = response.json() + assert data["total"] == 5 + assert data["has_more"] is True + assert len(data["results"]) == 3 @pytest.mark.asyncio @@ -192,7 +243,7 @@ async def test_search_with_type_filter( # Search with type filter response = await client.post( - f"{v2_project_url}/search/", json={"search_text": "Type", "note_types": ["note"]} + f"{v2_project_url}/search/", json={"text": "Type", "note_types": ["note"]} ) assert response.status_code == 200 @@ -225,7 +276,7 @@ async def test_search_with_date_filter( # Search with date filter response = await client.post( f"{v2_project_url}/search/", - json={"search_text": "Date Filtered", "after_date": "2024-01-01T00:00:00Z"}, + json={"text": "Date Filtered", "after_date": "2024-01-01T00:00:00Z"}, ) assert response.status_code == 200 @@ -246,12 +297,42 @@ async def test_search_empty_query( assert response.status_code in [200, 422] +@pytest.mark.asyncio +async def test_search_whitespace_text_is_treated_as_empty( + client: AsyncClient, + test_project: Project, + v2_project_url: str, + entity_repository, + search_service, + file_service, +): + """Whitespace-only text should not become an unfiltered project-wide search.""" + entity_data = { + "title": "Whitespace Regression Entity", + "note_type": "note", + "content_type": "text/markdown", + "file_path": "whitespace_regression.md", + "checksum": "whitespace123", + } + await create_test_entity( + test_project, entity_data, entity_repository, search_service, file_service + ) + + response = await client.post(f"{v2_project_url}/search/", json={"text": " "}) + + assert response.status_code == 200 + data = response.json() + assert data["total"] == 0 + assert data["has_more"] is False + assert data["results"] == [] + + @pytest.mark.asyncio async def test_search_invalid_project_id( client: AsyncClient, ): """Test searching with invalid project ID returns 404.""" - response = await client.post("/v2/projects/999999/search/", json={"search_text": "test"}) + response = await client.post("/v2/projects/999999/search/", json={"text": "test"}) assert response.status_code == 404 @@ -291,7 +372,7 @@ async def test_v2_search_endpoints_use_project_id_not_name( ): """Test that v2 search endpoints reject string project names.""" # Try to use project name instead of ID - should fail - response = await client.post(f"/v2/{test_project.name}/search/", json={"search_text": "test"}) + response = await client.post(f"/v2/{test_project.name}/search/", json={"text": "test"}) # FastAPI path validation should reject non-integer project_id assert response.status_code in [404, 422] @@ -307,11 +388,14 @@ async def test_search_router_returns_400_for_semantic_disabled( async def search(self, *args, **kwargs): raise SemanticSearchDisabledError("Semantic search is disabled for this project.") + async def count(self, *args, **kwargs): + raise SemanticSearchDisabledError("Semantic search is disabled for this project.") + app.dependency_overrides[get_search_service_v2_external] = lambda: RaisingSearchService() try: response = await client.post( f"{v2_project_url}/search/", - json={"search_text": "semantic query", "retrieval_mode": "vector"}, + json={"text": "semantic query", "retrieval_mode": "vector"}, ) finally: app.dependency_overrides.pop(get_search_service_v2_external, None) @@ -330,11 +414,14 @@ async def test_search_router_returns_400_for_semantic_missing_deps( async def search(self, *args, **kwargs): raise SemanticDependenciesMissingError("Semantic dependencies are missing.") + async def count(self, *args, **kwargs): + raise SemanticDependenciesMissingError("Semantic dependencies are missing.") + app.dependency_overrides[get_search_service_v2_external] = lambda: RaisingSearchService() try: response = await client.post( f"{v2_project_url}/search/", - json={"search_text": "semantic query", "retrieval_mode": "hybrid"}, + json={"text": "semantic query", "retrieval_mode": "hybrid"}, ) finally: app.dependency_overrides.pop(get_search_service_v2_external, None) @@ -353,6 +440,9 @@ async def test_search_router_returns_400_for_invalid_vector_query( async def search(self, *args, **kwargs): raise ValueError("Vector retrieval requires a text query.") + async def count(self, *args, **kwargs): + raise ValueError("Vector retrieval requires a text query.") + app.dependency_overrides[get_search_service_v2_external] = lambda: RaisingSearchService() try: response = await client.post( @@ -366,6 +456,56 @@ async def test_search_router_returns_400_for_invalid_vector_query( assert "Vector retrieval requires a text query" in response.json()["detail"] +@pytest.mark.asyncio +async def test_semantic_search_uses_probe_pagination_without_count( + client: AsyncClient, + app, + v2_project_url: str, +): + """Semantic searches should not run an extra count query.""" + now = datetime.now(timezone.utc) + fake_rows = [ + SearchIndexRow( + project_id=1, + id=row_id, + type="entity", + file_path=f"notes/semantic-{row_id}.md", + created_at=now, + updated_at=now, + title=f"Semantic Result {row_id}", + permalink=f"notes/semantic-{row_id}", + score=1.0 - (row_id / 10), + ) + for row_id in range(1, 4) + ] + + class FakeSearchService: + async def search(self, query, *, limit, offset): + assert query.retrieval_mode.value == "vector" + assert limit == 3 + assert offset == 0 + return fake_rows + + async def count(self, *args, **kwargs): + raise AssertionError("semantic search must not run count") + + app.dependency_overrides[get_search_service_v2_external] = lambda: FakeSearchService() + try: + response = await client.post( + f"{v2_project_url}/search/", + json={"text": "semantic query", "retrieval_mode": "vector"}, + params={"page": 1, "page_size": 2}, + ) + finally: + app.dependency_overrides.pop(get_search_service_v2_external, None) + + assert response.status_code == 200 + data = response.json() + assert data["total"] == 0 + assert data["has_more"] is True + assert len(data["results"]) == 2 + + @pytest.mark.asyncio async def test_search_has_more_when_more_results_exist( client: AsyncClient, @@ -459,11 +599,14 @@ async def test_search_result_includes_matched_chunk( async def search(self, *args, **kwargs): return [fake_row] + async def count(self, *args, **kwargs): + return 1 + app.dependency_overrides[get_search_service_v2_external] = lambda: FakeSearchService() try: response = await client.post( f"{v2_project_url}/search/", - json={"search_text": "pricing"}, + json={"text": "pricing"}, ) finally: app.dependency_overrides.pop(get_search_service_v2_external, None) @@ -500,11 +643,14 @@ async def test_search_result_omits_matched_chunk_when_none( async def search(self, *args, **kwargs): return [fake_row] + async def count(self, *args, **kwargs): + return 1 + app.dependency_overrides[get_search_service_v2_external] = lambda: FakeSearchService() try: response = await client.post( f"{v2_project_url}/search/", - json={"search_text": "general"}, + json={"text": "general"}, ) finally: app.dependency_overrides.pop(get_search_service_v2_external, None) diff --git a/tests/api/v2/test_search_router_telemetry.py b/tests/api/v2/test_search_router_telemetry.py index 92893843..5809402f 100644 --- a/tests/api/v2/test_search_router_telemetry.py +++ b/tests/api/v2/test_search_router_telemetry.py @@ -23,6 +23,9 @@ async def test_search_router_wraps_request_in_manual_operation(monkeypatch) -> N async def search(self, query, *, limit, offset): return [] + async def count(self, query): + return 0 + @contextmanager def fake_span(name: str, **attrs): operations.append((name, attrs)) diff --git a/tests/repository/test_postgres_search_repository.py b/tests/repository/test_postgres_search_repository.py index 20f455b1..3e733029 100644 --- a/tests/repository/test_postgres_search_repository.py +++ b/tests/repository/test_postgres_search_repository.py @@ -235,6 +235,7 @@ async def test_postgres_search_repository_tsquery_syntax_error_returns_empty( # Trailing boolean operator creates an invalid tsquery; repository should return [] results = await repo.search(search_text="coffee AND") assert results == [] + assert await repo.count(search_text="coffee AND") == 0 @pytest.mark.asyncio diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index 6998d199..ba4e08ef 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -318,6 +318,25 @@ async def test_sqlite_hybrid_search_raises_disabled_error(search_repository): ) +@pytest.mark.asyncio +@pytest.mark.parametrize("retrieval_mode", [SearchRetrievalMode.VECTOR, SearchRetrievalMode.HYBRID]) +async def test_count_rejects_semantic_modes_without_running_search(monkeypatch, retrieval_mode): + """Semantic counts must not materialize vector or hybrid retrieval.""" + repo = _ConcreteRepo() + search_calls = [] + + async def fail_if_search_runs(**kwargs): + search_calls.append(kwargs) + return [] + + monkeypatch.setattr(repo, "search", fail_if_search_runs) + + with pytest.raises(ValueError, match="Exact counts are only supported for full-text search"): + await repo.count(search_text="semantic query", retrieval_mode=retrieval_mode) + + assert search_calls == [] + + @pytest.mark.asyncio async def test_sync_entity_vectors_batch_flushes_at_configured_threshold(monkeypatch): """Batch sync should flush queued jobs at semantic_embedding_sync_batch_size boundaries.""" diff --git a/tests/schemas/test_search.py b/tests/schemas/test_search.py index b2fd8639..4c9fd71b 100644 --- a/tests/schemas/test_search.py +++ b/tests/schemas/test_search.py @@ -127,6 +127,7 @@ def test_search_response(): metadata={}, ), ] - response = SearchResponse(results=results, current_page=1, page_size=1) + response = SearchResponse(results=results, current_page=1, page_size=1, total=2) assert len(response.results) == 2 assert response.results[0].score > response.results[1].score + assert response.total == 2