From 02f8e866923d5793d2620076c709c920d99f2c4f Mon Sep 17 00:00:00 2001 From: phernandez Date: Tue, 18 Feb 2025 20:42:19 -0600 Subject: [PATCH] feat: add pagination to read_notes --- .../api/routers/resource_router.py | 8 +- src/basic_memory/mcp/tools/notes.py | 8 +- src/basic_memory/mcp/tools/search.py | 2 +- tests/api/test_resource_router.py | 77 +++++++++++++++++++ tests/mcp/test_tool_notes.py | 28 +++++++ 5 files changed, 119 insertions(+), 4 deletions(-) diff --git a/src/basic_memory/api/routers/resource_router.py b/src/basic_memory/api/routers/resource_router.py index 7fd3b404..3bc1c4bb 100644 --- a/src/basic_memory/api/routers/resource_router.py +++ b/src/basic_memory/api/routers/resource_router.py @@ -44,6 +44,8 @@ async def get_resource_content( file_service: FileServiceDep, background_tasks: BackgroundTasks, identifier: str, + page: int = 1, + page_size: int = 10, ) -> FileResponse: """Get resource content by identifier: name or permalink.""" logger.debug(f"Getting content for: {identifier}") @@ -52,6 +54,10 @@ async def get_resource_content( entity = await link_resolver.resolve_link(identifier) results = [entity] if entity else [] + # pagination for multiple results + limit = page_size + offset = (page - 1) * page_size + # search using the identifier as a permalink if not results: # if the identifier contains a wildcard, use GLOB search @@ -60,7 +66,7 @@ async def get_resource_content( if "*" in identifier else SearchQuery(permalink=identifier) ) - search_results = await search_service.search(query) + search_results = await search_service.search(query, limit, offset) if not search_results: raise HTTPException(status_code=404, detail=f"Resource not found: {identifier}") diff --git a/src/basic_memory/mcp/tools/notes.py b/src/basic_memory/mcp/tools/notes.py index 025c348d..7711d042 100644 --- a/src/basic_memory/mcp/tools/notes.py +++ b/src/basic_memory/mcp/tools/notes.py @@ -116,7 +116,7 @@ async def write_note( @mcp.tool(description="Read note content by title, permalink, relation, or pattern") -async def read_note(identifier: str) -> str: +async def read_note(identifier: str, page: int = 1, page_size: int = 10) -> str: """Get note content in unified diff format. The content is returned in a unified diff inspired format: @@ -134,6 +134,8 @@ async def read_note(identifier: str) -> str: - Note permalink ("docs/example") - Relation path ("docs/example/depends-on/other-doc") - Pattern match ("docs/*-architecture") + page: the page number of results to return (default 1) + page_size: the number of results to return per page (default 10) Returns: Document content in unified diff format. For single documents, returns @@ -171,7 +173,9 @@ async def read_note(identifier: str) -> str: with logfire.span("Reading note", identifier=identifier): # pyright: ignore [reportGeneralTypeIssues] logger.info(f"Reading note {identifier}") url = memory_url_path(identifier) - response = await call_get(client, f"/resource/{url}") + response = await call_get( + client, f"/resource/{url}", params={"page": page, "page_size": page_size} + ) return response.text diff --git a/src/basic_memory/mcp/tools/search.py b/src/basic_memory/mcp/tools/search.py index 035edcdf..b3445db7 100644 --- a/src/basic_memory/mcp/tools/search.py +++ b/src/basic_memory/mcp/tools/search.py @@ -27,7 +27,7 @@ async def search(query: SearchQuery, page: int = 1, page_size: int = 10) -> Sear Returns: SearchResponse with search results and metadata """ - with logfire.span("Searching for {query}", qurey=query): # pyright: ignore [reportGeneralTypeIssues] + with logfire.span("Searching for {query}", query=query): # pyright: ignore [reportGeneralTypeIssues] logger.info(f"Searching for {query}") response = await call_post( client, f"/search/?page={page}&page_size={page_size}", json=query.model_dump() diff --git a/tests/api/test_resource_router.py b/tests/api/test_resource_router.py index 4daa823a..aaa22085 100644 --- a/tests/api/test_resource_router.py +++ b/tests/api/test_resource_router.py @@ -37,6 +37,35 @@ async def test_get_resource_content(client, test_config, entity_repository): assert response.text == content +@pytest.mark.asyncio +async def test_get_resource_pagination(client, test_config, entity_repository): + """Test getting content by permalink with pagination.""" + # Create a test file + content = "# Test Content\n\nThis is a test file." + test_file = Path(test_config.home) / "test" / "test.md" + test_file.parent.mkdir(parents=True, exist_ok=True) + test_file.write_text(content) + + # Create entity referencing the file + entity = await entity_repository.create( + { + "title": "Test Entity", + "entity_type": "test", + "permalink": "test/test", + "file_path": "test/test.md", # Relative to config.home + "content_type": "text/markdown", + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + } + ) + + # Test getting the content + response = await client.get(f"/resource/{entity.permalink}", params={"page": 1, "page_size": 1}) + assert response.status_code == 200 + assert response.headers["content-type"] == "text/markdown; charset=utf-8" + assert response.text == content + + @pytest.mark.asyncio async def test_get_resource_by_title(client, test_config, entity_repository): """Test getting content by permalink.""" @@ -179,6 +208,54 @@ async def test_get_resource_entities(client, test_config, entity_repository): ) +@pytest.mark.asyncio +async def test_get_resource_entities_pagination(client, test_config, entity_repository): + """Test getting content by permalink match.""" + # Create entity + content1 = "# Test Content\n" + data = { + "title": "Test Entity", + "folder": "test", + "entity_type": "test", + "content": f"{content1}", + } + response = await client.post("/knowledge/entities", json=data) + entity_response = response.json() + entity1 = EntityResponse(**entity_response) + assert entity1 + + content2 = "# Related Content\n- links to [[Test Entity]]" + data = { + "title": "Related Entity", + "folder": "test", + "entity_type": "test", + "content": f"{content2}", + } + response = await client.post("/knowledge/entities", json=data) + entity_response = response.json() + entity2 = EntityResponse(**entity_response) + + assert len(entity2.relations) == 1 + + # Test getting second result + response = await client.get("/resource/test/*", params={"page": 2, "page_size": 1}) + assert response.status_code == 200 + assert response.headers["content-type"] == "text/markdown; charset=utf-8" + assert ( + """ +--- +title: Related Entity +type: test +permalink: test/related-entity +--- + +# Related Content +- links to [[Test Entity]] +""".strip() + in response.text + ) + + @pytest.mark.asyncio async def test_get_resource_relation(client, test_config, entity_repository): """Test getting content by relation permalink.""" diff --git a/tests/mcp/test_tool_notes.py b/tests/mcp/test_tool_notes.py index 24be8faf..c40c911b 100644 --- a/tests/mcp/test_tool_notes.py +++ b/tests/mcp/test_tool_notes.py @@ -221,6 +221,34 @@ async def test_multiple_notes(app): assert "--- memory://test/note-3" in result assert "Content 3" in result +@pytest.mark.asyncio +async def test_multiple_notes_pagination(app): + """Test creating and managing multiple notes.""" + # Create several notes + notes_data = [ + ("test/note-1", "Note 1", "test", "Content 1", ["tag1"]), + ("test/note-2", "Note 2", "test", "Content 2", ["tag1", "tag2"]), + ("test/note-3", "Note 3", "test", "Content 3", []), + ] + + for _, title, folder, content, tags in notes_data: + await notes.write_note(title=title, folder=folder, content=content, tags=tags) + + # Should be able to read each one + for permalink, title, folder, content, _ in notes_data: + note = await notes.read_note(permalink) + assert content in note + + # read multiple notes at once with pagination + result = await notes.read_note("test/*", page=1, page_size=2) + + # note we can't compare times + assert "--- memory://test/note-1" in result + assert "Content 1" in result + + assert "--- memory://test/note-2" in result + assert "Content 2" in result + @pytest.mark.asyncio async def test_delete_note_existing(app):