From 00d23a5ee15ddac4ea45e702dcd02ab9f0509276 Mon Sep 17 00:00:00 2001 From: phernandez Date: Tue, 18 Feb 2025 19:56:20 -0600 Subject: [PATCH] fix: add logfire spans to cli --- src/basic_memory/api/routers/search_router.py | 26 ++++++++++++------- src/basic_memory/mcp/tools/search.py | 8 ++++-- .../repository/search_repository.py | 3 +++ src/basic_memory/schemas/request.py | 2 +- src/basic_memory/schemas/search.py | 2 ++ src/basic_memory/services/context_service.py | 10 ++++--- src/basic_memory/services/search_service.py | 4 ++- tests/api/test_search_router.py | 12 +++++++++ tests/mcp/test_tool_search.py | 21 +++++++++++++++ tests/services/test_search_service.py | 19 ++++++++++++++ 10 files changed, 91 insertions(+), 16 deletions(-) diff --git a/src/basic_memory/api/routers/search_router.py b/src/basic_memory/api/routers/search_router.py index 6f0f2ce8..9acf7e00 100644 --- a/src/basic_memory/api/routers/search_router.py +++ b/src/basic_memory/api/routers/search_router.py @@ -2,27 +2,35 @@ from dataclasses import asdict -from fastapi import APIRouter, Depends, BackgroundTasks +from fastapi import APIRouter, BackgroundTasks -from basic_memory.services.search_service import SearchService from basic_memory.schemas.search import SearchQuery, SearchResult, SearchResponse -from basic_memory.deps import get_search_service +from basic_memory.deps import SearchServiceDep router = APIRouter(prefix="/search", tags=["search"]) @router.post("/", response_model=SearchResponse) -async def search(query: SearchQuery, search_service: SearchService = Depends(get_search_service)): +async def search( + query: SearchQuery, + search_service: SearchServiceDep, + page: int = 1, + page_size: int = 10, +): """Search across all knowledge and documents.""" - results = await search_service.search(query) + limit = page_size + offset = (page - 1) * page_size + results = await search_service.search(query, limit=limit, offset=offset) search_results = [SearchResult.model_validate(asdict(r)) for r in results] - return SearchResponse(results=search_results) + return SearchResponse( + results=search_results, + current_page=page, + page_size=page_size, + ) @router.post("/reindex") -async def reindex( - background_tasks: BackgroundTasks, search_service: SearchService = Depends(get_search_service) -): +async def reindex(background_tasks: BackgroundTasks, search_service: SearchServiceDep): """Recreate and populate the search index.""" await search_service.reindex_all(background_tasks=background_tasks) return {"status": "ok", "message": "Reindex initiated"} diff --git a/src/basic_memory/mcp/tools/search.py b/src/basic_memory/mcp/tools/search.py index c532c07f..035edcdf 100644 --- a/src/basic_memory/mcp/tools/search.py +++ b/src/basic_memory/mcp/tools/search.py @@ -12,7 +12,7 @@ from basic_memory.mcp.async_client import client @mcp.tool( description="Search across all content in basic-memory, including documents and entities", ) -async def search(query: SearchQuery) -> SearchResponse: +async def search(query: SearchQuery, page: int = 1, page_size: int = 10) -> SearchResponse: """Search across all content in basic-memory. Args: @@ -21,11 +21,15 @@ async def search(query: SearchQuery) -> SearchResponse: - types: Optional list of content types to search ("document" or "entity") - entity_types: Optional list of entity types to filter by - after_date: Optional date filter for recent content + page: the page number of results to return (default 1) + page_size: the number of results to return per page (default 10) Returns: SearchResponse with search results and metadata """ with logfire.span("Searching for {query}", qurey=query): # pyright: ignore [reportGeneralTypeIssues] logger.info(f"Searching for {query}") - response = await call_post(client, "/search/", json=query.model_dump()) + response = await call_post( + client, f"/search/?page={page}&page_size={page_size}", json=query.model_dump() + ) return SearchResponse.model_validate(response.json()) diff --git a/src/basic_memory/repository/search_repository.py b/src/basic_memory/repository/search_repository.py index 6f3a51c1..de87dab2 100644 --- a/src/basic_memory/repository/search_repository.py +++ b/src/basic_memory/repository/search_repository.py @@ -114,6 +114,7 @@ class SearchRepository: after_date: Optional[datetime] = None, entity_types: Optional[List[str]] = None, limit: int = 10, + offset: int = 0, ) -> List[SearchIndexRow]: """Search across all indexed content with fuzzy matching.""" conditions = [] @@ -169,6 +170,7 @@ class SearchRepository: # set limit on search query params["limit"] = limit + params["offset"] = offset # Build WHERE clause where_clause = " AND ".join(conditions) if conditions else "1=1" @@ -194,6 +196,7 @@ class SearchRepository: WHERE {where_clause} ORDER BY score ASC {order_by_clause} LIMIT :limit + OFFSET :offset """ logger.debug(f"Search {sql} params: {params}") diff --git a/src/basic_memory/schemas/request.py b/src/basic_memory/schemas/request.py index d8803d85..3fa7a80c 100644 --- a/src/basic_memory/schemas/request.py +++ b/src/basic_memory/schemas/request.py @@ -51,7 +51,7 @@ class GetEntitiesRequest(BaseModel): discovered through search. """ - permalinks: Annotated[List[Permalink], MinLen(1)] + permalinks: Annotated[List[Permalink], MinLen(1), MaxLen(10)] class CreateRelationsRequest(BaseModel): diff --git a/src/basic_memory/schemas/search.py b/src/basic_memory/schemas/search.py index 87a6c1ec..f43a8ebd 100644 --- a/src/basic_memory/schemas/search.py +++ b/src/basic_memory/schemas/search.py @@ -102,6 +102,8 @@ class SearchResponse(BaseModel): """Wrapper for search results.""" results: List[SearchResult] + current_page: int + page_size: int # Schema for future advanced search endpoint diff --git a/src/basic_memory/services/context_service.py b/src/basic_memory/services/context_service.py index b8254e92..58901c07 100644 --- a/src/basic_memory/services/context_service.py +++ b/src/basic_memory/services/context_service.py @@ -66,15 +66,19 @@ class ContextService: # Pattern matching - use search if "*" in path: logger.debug(f"Pattern search for '{path}'") - primary = await self.search_repository.search(permalink_match=path) + primary = await self.search_repository.search( + permalink_match=path, limit=max_results + ) # Direct lookup for exact path else: logger.debug(f"Direct lookup for '{path}'") - primary = await self.search_repository.search(permalink=path) + primary = await self.search_repository.search(permalink=path, limit=max_results) else: logger.debug(f"Build context for '{types}'") - primary = await self.search_repository.search(types=types, after_date=since) + primary = await self.search_repository.search( + types=types, after_date=since, limit=max_results + ) # Get type_id pairs for traversal diff --git a/src/basic_memory/services/search_service.py b/src/basic_memory/services/search_service.py index 53fa08f0..738f3ba3 100644 --- a/src/basic_memory/services/search_service.py +++ b/src/basic_memory/services/search_service.py @@ -51,7 +51,7 @@ class SearchService: logger.info("Reindex complete") - async def search(self, query: SearchQuery) -> List[SearchIndexRow]: + async def search(self, query: SearchQuery, limit=10, offset=0) -> List[SearchIndexRow]: """Search across all indexed content. Supports three modes: @@ -84,6 +84,8 @@ class SearchService: types=query.types, entity_types=query.entity_types, after_date=after_date, + limit=limit, + offset=offset, ) return results diff --git a/tests/api/test_search_router.py b/tests/api/test_search_router.py index e681025b..cd390f71 100644 --- a/tests/api/test_search_router.py +++ b/tests/api/test_search_router.py @@ -35,6 +35,18 @@ async def test_search_basic(client, indexed_entity): assert found, "Expected to find indexed entity in results" +@pytest.mark.asyncio +async def test_search_basic_pagination(client, indexed_entity): + """Test basic text search.""" + response = await client.post("/search/?page=3&page_size=1", json={"text": "search"}) + assert response.status_code == 200 + search_results = SearchResponse.model_validate(response.json()) + assert len(search_results.results) == 1 + + assert search_results.current_page == 3 + assert search_results.page_size == 1 + + @pytest.mark.asyncio async def test_search_with_type_filter(client, indexed_entity): """Test search with type filter.""" diff --git a/tests/mcp/test_tool_search.py b/tests/mcp/test_tool_search.py index a2b398be..e88fb873 100644 --- a/tests/mcp/test_tool_search.py +++ b/tests/mcp/test_tool_search.py @@ -29,6 +29,27 @@ async def test_search_basic(client): assert any(r.permalink == "test/test-search-note" for r in response.results) +@pytest.mark.asyncio +async def test_search_pagination(client): + """Test basic search functionality.""" + # Create a test note + result = await notes.write_note( + title="Test Search Note", + folder="test", + content="# Test\nThis is a searchable test note", + tags=["test", "search"], + ) + assert result + + # Search for it + query = SearchQuery(text="searchable") + response = await search(query, page=1, page_size=1) + + # Verify results + assert len(response.results) == 1 + assert any(r.permalink == "test/test-search-note" for r in response.results) + + @pytest.mark.asyncio async def test_search_with_type_filter(client): """Test search with entity type filter.""" diff --git a/tests/services/test_search_service.py b/tests/services/test_search_service.py index 917c30a2..8e0d1bb5 100644 --- a/tests/services/test_search_service.py +++ b/tests/services/test_search_service.py @@ -19,6 +19,25 @@ async def test_search_permalink(search_service, test_graph): assert "test/root" in r.permalink +@pytest.mark.asyncio +async def test_search_limit_offset(search_service, test_graph): + """Exact permalink""" + results = await search_service.search(SearchQuery(permalink_match="test/*")) + assert len(results) > 1 + + results = await search_service.search(SearchQuery(permalink_match="test/*"), limit=1) + assert len(results) == 1 + + results = await search_service.search(SearchQuery(permalink_match="test/*"), limit=100) + num_results = len(results) + + # assert offset + offset_results = await search_service.search( + SearchQuery(permalink_match="test/*"), limit=100, offset=1 + ) + assert len(offset_results) == num_results - 1 + + @pytest.mark.asyncio async def test_search_permalink_observations_wildcard(search_service, test_graph): """Pattern matching"""