From 0a177f17c8316ad53256f64c50172137c910ec9c Mon Sep 17 00:00:00 2001 From: phernandez Date: Mon, 23 Dec 2024 17:26:18 -0600 Subject: [PATCH] add /documents endpoints to mcp server --- src/basic_memory/api/routers/documents.py | 26 +--- src/basic_memory/mcp/server.py | 169 ++++++++++++++++++---- src/basic_memory/schemas/__init__.py | 4 + src/basic_memory/schemas/request.py | 6 +- tests/mcp/test_list_tools.py | 17 +++ tests/mcp/test_mcp_server.py | 85 +++++++++++ 6 files changed, 260 insertions(+), 47 deletions(-) diff --git a/src/basic_memory/api/routers/documents.py b/src/basic_memory/api/routers/documents.py index b2613176..9744bd06 100644 --- a/src/basic_memory/api/routers/documents.py +++ b/src/basic_memory/api/routers/documents.py @@ -5,7 +5,7 @@ from typing import List from fastapi import APIRouter, HTTPException from basic_memory.deps import DocumentServiceDep -from basic_memory.schemas.request import DocumentCreate, DocumentUpdate, DocumentPatch +from basic_memory.schemas.request import DocumentCreateRequest, DocumentUpdateRequest from basic_memory.schemas.response import DocumentResponse, DocumentCreateResponse from basic_memory.services.document_service import ( DocumentNotFoundError, @@ -18,11 +18,11 @@ router = APIRouter(prefix="/documents", tags=["documents"]) @router.post("/", response_model=DocumentCreateResponse, status_code=201) async def create_document( - doc: DocumentCreate, + doc: DocumentCreateRequest, service: DocumentServiceDep, ) -> DocumentCreateResponse: """Create a new document. - + The document will be created with appropriate frontmatter including: - Generated ID - Creation timestamp @@ -61,10 +61,7 @@ async def get_document( response = DocumentResponse.model_validate(doc_dict) return response except DocumentNotFoundError: - raise HTTPException( - status_code=404, - detail=f"Document not found: {id}" - ) + raise HTTPException(status_code=404, detail=f"Document not found: {id}") except DocumentWriteError as e: raise HTTPException(status_code=400, detail=str(e)) @@ -72,15 +69,14 @@ async def get_document( @router.put("/{id:int}", response_model=DocumentResponse) async def update_document( id: int, - doc: DocumentUpdate, + doc: DocumentUpdateRequest, service: DocumentServiceDep, ) -> DocumentResponse: """Update a document by ID.""" # Verify IDs match if doc.id != id: raise HTTPException( - status_code=400, - detail="Document ID in URL must match ID in request body" + status_code=400, detail="Document ID in URL must match ID in request body" ) try: @@ -92,10 +88,7 @@ async def update_document( doc_dict = document.__dict__ | {"content": doc.content} return DocumentResponse.model_validate(doc_dict) except DocumentNotFoundError: - raise HTTPException( - status_code=404, - detail=f"Document not found: {id}" - ) + raise HTTPException(status_code=404, detail=f"Document not found: {id}") except DocumentWriteError as e: raise HTTPException(status_code=400, detail=str(e)) @@ -109,9 +102,6 @@ async def delete_document( try: await service.delete_document_by_id(id) except DocumentNotFoundError: - raise HTTPException( - status_code=404, - detail=f"Document not found: {id}" - ) + raise HTTPException(status_code=404, detail=f"Document not found: {id}") except DocumentWriteError as e: raise HTTPException(status_code=400, detail=str(e)) diff --git a/src/basic_memory/mcp/server.py b/src/basic_memory/mcp/server.py index a6b250b7..450a0f01 100644 --- a/src/basic_memory/mcp/server.py +++ b/src/basic_memory/mcp/server.py @@ -24,7 +24,7 @@ from mcp.types import ( ServerCapabilities, ToolsCapability, ) -from pydantic import TypeAdapter, AnyUrl +from pydantic import TypeAdapter, AnyUrl, ValidationError from basic_memory.api.app import app as fastapi_app from basic_memory.schemas import ( @@ -36,6 +36,8 @@ from basic_memory.schemas import ( DeleteEntitiesRequest, DeleteObservationsRequest, DeleteRelationsRequest, + DocumentCreateRequest, + DocumentUpdateRequest, ) BASE_URL = "http://test" @@ -54,6 +56,7 @@ async def handle_list_tools() -> List[Tool]: """Define the available tools.""" logger.debug("Listing available tools") return [ + # Knowledge graph tools Tool( name="create_entities", description="Create multiple new entities", @@ -94,18 +97,83 @@ async def handle_list_tools() -> List[Tool]: description="Delete relations", inputSchema=DeleteRelationsRequest.model_json_schema(), ), + # Document tools + Tool( + name="create_document", + description="Create a new document", + inputSchema=DocumentCreateRequest.model_json_schema(), + ), + Tool( + name="list_documents", + description="List all documents", + inputSchema={ + "type": "object", + "properties": {}, + "additionalProperties": False, + }, + ), + Tool( + name="get_document", + description="Get a document by ID", + inputSchema={ + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 1, + "description": "Document ID" + }, + }, + "required": ["id"], + "additionalProperties": False, + }, + ), + Tool( + name="update_document", + description="Update a document by ID", + inputSchema=DocumentUpdateRequest.model_json_schema(), + ), + Tool( + name="delete_document", + description="Delete a document by ID", + inputSchema={ + "type": "object", + "properties": { + "id": { + "type": "integer", + "minimum": 1, + "description": "Document ID" + }, + }, + "required": ["id"], + "additionalProperties": False, + }, + ), ] -async def call_tool_endpoint(endpoint: str, json: dict[str, Any]): +async def call_tool_endpoint(endpoint: str, json: dict[str, Any], method: str = "post"): """Makes a request to a FastAPI endpoint with a fresh client.""" async with AsyncClient( transport=ASGITransport(app=fastapi_app), base_url=BASE_URL, timeout=30.0 ) as client: - logger.debug(f"Calling API endpoint {endpoint} with arguments: {json}") - response = await client.post(endpoint, json=json) - logger.debug(response.json()) - return response + logger.debug(f"Calling API endpoint {endpoint} with {method}: {json}") + try: + if method == "post": + response = await client.post(endpoint, json=json) + elif method == "get": + # For GET requests, don't send body + params = {k:v for k,v in json.items() if k != "id"} + response = await client.get(endpoint, params=params) + elif method == "put": + response = await client.put(endpoint, json=json) + elif method == "delete": + response = await client.delete(endpoint) + logger.debug(response.json() if response.content else "No content") + return response + except ValidationError as e: + logger.error(f"Validation error: {e}") + raise McpError(INVALID_PARAMS, str(e)) def create_response(data: Dict[str, Any]) -> EmbeddedResource: @@ -128,44 +196,93 @@ async def handle_call_tool(name: str, arguments: Dict[str, Any]) -> List[Embedde logger.info(f"Tool call: {name}") logger.debug(f"Arguments: {arguments}") - # Map tools to FastAPI endpoints + # Map tools to FastAPI endpoints and methods handlers = { - "create_entities": "/knowledge/entities", - "search_nodes": "/knowledge/search", - "open_nodes": "/knowledge/nodes", - "add_observations": "/knowledge/observations", - "create_relations": "/knowledge/relations", - "delete_entities": "/knowledge/entities/delete", - "delete_observations": "/knowledge/observations/delete", - "delete_relations": "/knowledge/relations/delete", + # Knowledge graph endpoints + "create_entities": ("/knowledge/entities", "post"), + "search_nodes": ("/knowledge/search", "post"), + "open_nodes": ("/knowledge/nodes", "post"), + "add_observations": ("/knowledge/observations", "post"), + "create_relations": ("/knowledge/relations", "post"), + "delete_entities": ("/knowledge/entities/delete", "post"), + "delete_observations": ("/knowledge/observations/delete", "post"), + "delete_relations": ("/knowledge/relations/delete", "post"), + # Document endpoints + "create_document": ("/documents", "post"), + "list_documents": ("/documents", "get"), + "get_document": ("/documents/{id}", "get"), + "update_document": ("/documents/{id}", "put"), + "delete_document": ("/documents/{id}", "delete"), } # Get handler for tool - endpoint = handlers.get(name) - if endpoint is None: + handler = handlers.get(name) + if handler is None: raise McpError(METHOD_NOT_FOUND, f"Unknown tool: {name}") + endpoint, method = handler + + # Validate tool arguments for common fields + if "id" in arguments: + if not isinstance(arguments["id"], int): + raise McpError(INVALID_PARAMS, "ID must be an integer") + if arguments["id"] < 1: + raise McpError(INVALID_PARAMS, "ID must be greater than 0") + + if name == "create_document" and ( + "path" not in arguments or + "content" not in arguments + ): + raise McpError(INVALID_PARAMS, "Document creation requires path and content") + + if name == "update_document" and ( + "content" not in arguments or + "id" not in arguments + ): + raise McpError(INVALID_PARAMS, "Document update requires content and ID") + + # Format endpoint for ID-based routes + if "{id}" in endpoint: + id = arguments.get("id") + if id is None: + raise McpError(INVALID_PARAMS, "ID parameter required") + endpoint = endpoint.format(id=id) + # Make API call - response = await call_tool_endpoint(endpoint, arguments) + response = await call_tool_endpoint(endpoint, arguments, method) # Handle HTTP errors if response.status_code >= 400: error_data = response.json() - if response.status_code == 404: - raise McpError(METHOD_NOT_FOUND, error_data.get("detail", "Not found")) - elif response.status_code == 422: - raise McpError(INVALID_PARAMS, error_data.get("detail", "Invalid parameters")) + error_detail = error_data.get("detail", "") + + # For validation errors, always raise INVALID_PARAMS + if response.status_code == 422: + raise McpError(INVALID_PARAMS, error_detail) + # For document not found, use INVALID_PARAMS + elif response.status_code == 404 and "Document not found" in str(error_detail): + raise McpError(INVALID_PARAMS, error_detail) + # For other 404s, use METHOD_NOT_FOUND + elif response.status_code == 404: + raise McpError(METHOD_NOT_FOUND, error_detail or "Not found") else: - raise McpError(INTERNAL_ERROR, error_data.get("detail", "Internal error")) + raise McpError(INTERNAL_ERROR, error_detail or "Internal error") # Create response - result = create_response(response.json()) + if response.content: # Some endpoints (like DELETE) return no content + result = create_response(response.json()) + else: + result = create_response({"status": "success"}) + logger.debug(f"Tool call successful: {result}") return [result] - except ValueError as e: + except ValidationError as e: logger.error(f"Validation error: {e}") raise McpError(INVALID_PARAMS, str(e)) + except ValueError as e: + logger.error(f"Value error: {e}") + raise McpError(INVALID_PARAMS, str(e)) except Exception as e: logger.error(f"Error handling tool call: {e}") if isinstance(e, McpError): @@ -240,4 +357,4 @@ if __name__ == "__main__": logger.info("Server stopped by user") except Exception as e: logger.error(f"Fatal server error: {e}") - sys.exit(1) + sys.exit(1) \ No newline at end of file diff --git a/src/basic_memory/schemas/__init__.py b/src/basic_memory/schemas/__init__.py index 0ecbe5f2..197eef8c 100644 --- a/src/basic_memory/schemas/__init__.py +++ b/src/basic_memory/schemas/__init__.py @@ -28,6 +28,8 @@ from basic_memory.schemas.request import ( SearchNodesRequest, OpenNodesRequest, CreateRelationsRequest, + DocumentCreateRequest, + DocumentUpdateRequest, ) # Response models @@ -57,6 +59,8 @@ __all__ = [ "SearchNodesRequest", "OpenNodesRequest", "CreateRelationsRequest", + "DocumentCreateRequest", + "DocumentUpdateRequest", # Responses "SQLAlchemyModel", "ObservationResponse", diff --git a/src/basic_memory/schemas/request.py b/src/basic_memory/schemas/request.py index 362480b3..e574607d 100644 --- a/src/basic_memory/schemas/request.py +++ b/src/basic_memory/schemas/request.py @@ -201,13 +201,13 @@ class CreateRelationsRequest(BaseModel): ## document -class DocumentCreate(BaseModel): +class DocumentCreateRequest(BaseModel): path: str content: str doc_metadata: Optional[Dict[str, Any]] = None -class DocumentUpdate(BaseModel): +class DocumentUpdateRequest(BaseModel): """Update an existing document by ID.""" id: int # Document ID is required for updates @@ -215,7 +215,7 @@ class DocumentUpdate(BaseModel): doc_metadata: Optional[Dict[str, Any]] = None -class DocumentPatch(BaseModel): +class DocumentPatchRequest(BaseModel): id: int content: Optional[str] = None doc_metadata: Optional[Dict[str, Any]] = None diff --git a/tests/mcp/test_list_tools.py b/tests/mcp/test_list_tools.py index b0ffe72a..038ce287 100644 --- a/tests/mcp/test_list_tools.py +++ b/tests/mcp/test_list_tools.py @@ -13,6 +13,7 @@ async def test_list_tools(app): # Check each expected tool is present expected_tools = { + # Knowledge graph tools "create_entities", "search_nodes", "open_nodes", @@ -21,6 +22,12 @@ async def test_list_tools(app): "delete_entities", "delete_observations", "delete_relations", + # Document tools + "create_document", + "list_documents", + "get_document", + "update_document", + "delete_document", } found_tools = {t.name: t for t in tools} @@ -30,3 +37,13 @@ async def test_list_tools(app): search_schema = found_tools["search_nodes"].inputSchema assert "query" in search_schema["properties"] assert search_schema["required"] == ["query"] + + # Verify document tool schemas + create_doc_schema = found_tools["create_document"].inputSchema + assert "path" in create_doc_schema["properties"] + assert "content" in create_doc_schema["properties"] + assert set(create_doc_schema["required"]) == {"path", "content"} + + get_doc_schema = found_tools["get_document"].inputSchema + assert "id" in get_doc_schema["properties"] + assert get_doc_schema["required"] == ["id"] \ No newline at end of file diff --git a/tests/mcp/test_mcp_server.py b/tests/mcp/test_mcp_server.py index dbf1c202..2dee0f1c 100644 --- a/tests/mcp/test_mcp_server.py +++ b/tests/mcp/test_mcp_server.py @@ -26,6 +26,14 @@ async def test_missing_required_field(app): await handle_call_tool("create_entities", {}) assert "entities" in str(exc.value).lower() + with pytest.raises(McpError) as exc: + await handle_call_tool("create_document", {}) + assert "path" in str(exc.value).lower() or "content" in str(exc.value).lower() + + with pytest.raises(McpError) as exc: + await handle_call_tool("get_document", {}) + assert "id" in str(exc.value).lower() + @pytest.mark.asyncio async def test_empty_arrays(app): @@ -50,6 +58,10 @@ async def test_invalid_field_types(app): await handle_call_tool("create_entities", {"entities": "not an array"}) assert "array" in str(exc.value).lower() or "list" in str(exc.value).lower() + with pytest.raises(McpError) as exc: + await handle_call_tool("get_document", {"id": "not an integer"}) + assert "integer" in str(exc.value).lower() + @pytest.mark.asyncio async def test_invalid_nested_fields(app): @@ -69,6 +81,17 @@ async def test_invalid_nested_fields(app): ) assert "entity_type" in str(exc.value).lower() + with pytest.raises(McpError) as exc: + await handle_call_tool( + "create_document", + { + "path": "test.md", + "content": "test", + "doc_metadata": "not an object" # Should be dict/null + }, + ) + assert "doc_metadata" in str(exc.value).lower() + @pytest.mark.asyncio async def test_invalid_relation_format_to_id(app): @@ -146,3 +169,65 @@ async def test_edge_case_validation_search_len(app): "search_nodes", {"query": "x" * 10000}, # Extremely long query ) + + +@pytest.mark.asyncio +async def test_document_endpoint_validation(app): + """Test validation specific to document endpoints.""" + # Invalid ID format for get_document + with pytest.raises(McpError) as exc: + await handle_call_tool("get_document", {"id": -1}) + assert INVALID_PARAMS == exc.value.args[0] + + # Invalid document path + with pytest.raises(McpError) as exc: + await handle_call_tool( + "create_document", + { + "path": "", # Empty path + "content": "test content" + } + ) + assert INVALID_PARAMS == exc.value.args[0] + + # Update without ID match + with pytest.raises(McpError) as exc: + await handle_call_tool( + "update_document", + { + "id": 1, + "content": "new content", + "doc_metadata": None + } + ) + assert "id" in str(exc.value).lower() + + +@pytest.mark.asyncio +async def test_document_http_methods(app): + """Test that document endpoints use correct HTTP methods.""" + # Test GET endpoints + await handle_call_tool("list_documents", {}) + await handle_call_tool("get_document", {"id": 1}) + + # Test POST endpoint + await handle_call_tool( + "create_document", + { + "path": "test.md", + "content": "test content" + } + ) + + # Test PUT endpoint + await handle_call_tool( + "update_document", + { + "id": 1, + "content": "updated content", + "doc_metadata": None + } + ) + + # Test DELETE endpoint + await handle_call_tool("delete_document", {"id": 1}) \ No newline at end of file