diff --git a/src/basic_memory/mcp/server.py b/src/basic_memory/mcp/server.py index 6e1920b8..24a3c910 100644 --- a/src/basic_memory/mcp/server.py +++ b/src/basic_memory/mcp/server.py @@ -1,77 +1,201 @@ """MCP server implementation for basic-memory.""" from pathlib import Path -from typing import Annotated, List, Dict, Any +from typing import List, Dict, Any from mcp.server import Server from mcp.types import Tool, TextContent, METHOD_NOT_FOUND, INVALID_PARAMS, INTERNAL_ERROR from mcp.shared.exceptions import McpError -from pydantic import BaseModel, Field from basic_memory import deps +from basic_memory.schemas import EntityIn, ObservationIn, RelationIn + +class MemoryServer(Server): + """Extended server class that exposes handlers for testing.""" + + def __init__(self): + super().__init__("basic-memory") + self.register_handlers() + + def register_handlers(self): + """Register all handlers with proper decorators.""" + + @self.list_tools() + async def handle_list_tools() -> List[Tool]: + """Define the available tools.""" + return [ + Tool( + name="create_entities", + description="Create multiple new entities in the knowledge graph", + inputSchema={ + "type": "object", + "properties": { + "entities": { + "type": "array", + "items": EntityIn.model_json_schema(), + "minItems": 1 + } + }, + "required": ["entities"] + } + ), + Tool( + name="search_nodes", + description="Search for nodes in the knowledge graph", + inputSchema={ + "type": "object", + "properties": { + "query": { + "type": "string", + "minLength": 1 + } + }, + "required": ["query"] + } + ), + Tool( + name="open_nodes", + description="Open specific nodes by their names", + inputSchema={ + "type": "object", + "properties": { + "names": { + "type": "array", + "items": {"type": "string"}, + "minItems": 1 + } + }, + "required": ["names"] + } + ), + Tool( + name="add_observations", + description="Add observations to existing entities", + inputSchema={ + "type": "object", + "properties": { + "entityId": {"type": "string"}, + "observations": { + "type": "array", + "items": ObservationIn.model_json_schema(), + "minItems": 1 + } + }, + "required": ["entityId", "observations"] + } + ), + Tool( + name="create_relations", + description="Create relations between entities", + inputSchema={ + "type": "object", + "properties": { + "relations": { + "type": "array", + "items": RelationIn.model_json_schema(), + "minItems": 1 + } + }, + "required": ["relations"] + } + ), + Tool( + name="delete_entities", + description="Delete entities from the knowledge graph", + inputSchema={ + "type": "object", + "properties": { + "names": { + "type": "array", + "items": {"type": "string"}, + "minItems": 1 + } + }, + "required": ["names"] + } + ), + Tool( + name="delete_observations", + description="Delete observations from entities", + inputSchema={ + "type": "object", + "properties": { + "deletions": { + "type": "array", + "items": { + "type": "object", + "properties": { + "entityName": {"type": "string"}, + "observations": { + "type": "array", + "items": {"type": "string"}, + "minItems": 1 + } + }, + "required": ["entityName", "observations"] + } + } + }, + "required": ["deletions"] + } + ) + ] + + @self.call_tool() + async def handle_call_tool(name: str, arguments: Dict[str, Any]) -> List[TextContent]: + """Handle tool calls by delegating to the memory service.""" + try: + # Get project path (could come from context in future) + project_path = Path.home() / ".basic-memory" / "projects" / "default" + + # Get services with proper lifecycle management + async with deps.get_project_services(project_path) as memory_service: + match name: + case "create_entities": + # Each entity in arguments["entities"] will be validated by EntityIn + result = await memory_service.create_entities(arguments["entities"]) + return [TextContent(type="text", text=str(result))] + + case "search_nodes": + result = await memory_service.search_nodes(arguments["query"]) + return [TextContent(type="text", text=str(result))] + + case "open_nodes": + result = await memory_service.open_nodes(arguments["names"]) + return [TextContent(type="text", text=str(result))] + + case "add_observations": + result = await memory_service.add_observations(arguments) + return [TextContent(type="text", text=str(result))] + + case "create_relations": + result = await memory_service.create_relations(arguments["relations"]) + return [TextContent(type="text", text=str(result))] + + case "delete_entities": + await memory_service.delete_entities(arguments["names"]) + return [TextContent(type="text", text="Entities deleted")] + + case "delete_observations": + await memory_service.delete_observations(arguments["deletions"]) + return [TextContent(type="text", text="Observations deleted")] + + case _: + raise McpError( + METHOD_NOT_FOUND, + f"Unknown tool: {name}" + ) + + except ValueError as e: + raise McpError(INVALID_PARAMS, str(e)) + except Exception as e: + raise McpError(INTERNAL_ERROR, str(e)) + + # Store handlers as instance attributes for testing + self.handle_list_tools = handle_list_tools + self.handle_call_tool = handle_call_tool # Create server instance -server = Server("basic-memory") - -# Parameter models -class CreateEntitiesParams(BaseModel): - """Parameters for creating entities.""" - entities: Annotated[List[Dict[str, Any]], Field( - description="List of entities to create", - min_items=1 - )] - -class SearchNodesParams(BaseModel): - """Parameters for searching nodes.""" - query: Annotated[str, Field( - description="Search query to match against entities" - )] - -@server.list_tools() -async def list_tools() -> List[Tool]: - """Define the available tools.""" - return [ - Tool( - name="create_entities", - description="Create multiple new entities in the knowledge graph", - inputSchema=CreateEntitiesParams.model_json_schema() - ), - Tool( - name="search_nodes", - description="Search for nodes in the knowledge graph", - inputSchema=SearchNodesParams.model_json_schema() - ) - ] - -@server.call_tool() -async def call_tool(name: str, arguments: Dict[str, Any]) -> List[TextContent]: - """Handle tool calls by delegating to the memory service.""" - try: - # Get project path (could come from context in future) - project_path = Path.home() / ".basic-memory" / "projects" / "default" - - # Get services with proper lifecycle management - async with deps.get_project_services(project_path) as memory_service: - match name: - case "create_entities": - params = CreateEntitiesParams(**arguments) - result = await memory_service.create_entities(params.entities) - return [TextContent(type="text", text=str(result))] - - case "search_nodes": - params = SearchNodesParams(**arguments) - result = await memory_service.search_nodes(params.query) - return [TextContent(type="text", text=str(result))] - - case _: - raise McpError( - METHOD_NOT_FOUND, - f"Unknown tool: {name}" - ) - - except ValueError as e: - raise McpError(INVALID_PARAMS, str(e)) - except Exception as e: - raise McpError(INTERNAL_ERROR, str(e)) +server = MemoryServer() async def run_server(): """Run the MCP server.""" diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index 7780edde..14c2c3dc 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -7,8 +7,7 @@ from datetime import datetime, UTC from typing import List, Optional, Dict, Any from uuid import uuid4 -from pydantic import BaseModel, model_validator - +from pydantic import BaseModel, Field, model_validator class ObservationIn(BaseModel): """Schema for creating a single observation.""" @@ -17,29 +16,37 @@ class ObservationIn(BaseModel): class ObservationsIn(BaseModel): """Schema for adding observations to an entity.""" - entity_id: str # Maps to Entity.id + entity_id: str = Field(alias="entityId") # Maps to Entity.id observations: List[ObservationIn] + class Config: + populate_by_name = True + class ObservationOut(ObservationIn): """Schema for observation data returned from the service.""" id: int class ObservationsOut(BaseModel): """Schema for bulk observation operation results.""" - entity_id: str + entity_id: str = Field(alias="entityId") observations: List[ObservationOut] + class Config: + populate_by_name = True class RelationIn(BaseModel): """ Represents a directed edge between entities in the knowledge graph. Relations are always stored in active voice (e.g. "created", "teaches", etc.) """ - from_id: str - to_id: str - relation_type: str + from_id: str = Field(alias="fromId") + to_id: str = Field(alias="toId") + relation_type: str = Field(alias="relationType") context: Optional[str] = None + class Config: + populate_by_name = True + class RelationOut(BaseModel): id: int @@ -47,7 +54,7 @@ class EntityBase(BaseModel): # id assigned at creation via model_validator id: str name: str - entity_type: str + entity_type: str = Field(alias="entityType") @model_validator(mode='before') @classmethod @@ -72,12 +79,13 @@ class EntityIn(EntityBase): observations: List[ObservationIn] = [] relations: List[RelationIn] = [] + class Config: + populate_by_name = True class EntityOut(EntityBase): - """ - Represents a node in our knowledge graph - could be a person, project, - concept, etc. Each entity has a unique name, a type, and a list of - associated observations. - """ + """Schema for entity data returned from the service.""" observations: List[ObservationOut] = [] relations: List[RelationOut] = [] + + class Config: + populate_by_name = True \ No newline at end of file diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 9adc204e..a6708d73 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -1,8 +1,7 @@ """Tests for the MCP server implementation.""" import pytest -from pathlib import Path -from mcp.types import Tool, TextContent +from mcp.types import TextContent from basic_memory.mcp import server @pytest.fixture @@ -11,49 +10,87 @@ def anyio_backend(): @pytest.fixture def test_entity_data(): + """Sample data for creating a test entity using camelCase (like MCP will).""" return { "entities": [{ - "name": "Test Entity", + "name": "Test Entity CamelCase", "entityType": "test", - "observations": ["This is a test observation"] + "observations": [{"content": "This is a test observation"}] + }] + } + +@pytest.fixture +def test_entity_snake_case(): + """Same test data but using snake_case to test schema flexibility.""" + return { + "entities": [{ + "name": "Test Entity SnakeCase", + "entity_type": "test", + "observations": [{"content": "This is a test observation"}] }] } @pytest.mark.anyio async def test_list_tools(): """Test that server exposes expected tools.""" - tools = server.list_tools() + tools = await server.handle_list_tools() - assert len(tools) == 2 # We have create_entities and search_nodes for now + # Check each expected tool is present + expected_tools = { + "create_entities", "search_nodes", "open_nodes", + "add_observations", "create_relations", + "delete_entities", "delete_observations" + } - # Verify create_entities tool - create_tool = next(t for t in tools if t.name == "create_entities") - assert create_tool.inputSchema["required"] == ["entities"] - assert "entities" in create_tool.inputSchema["properties"] + found_tools = {t.name: t for t in tools} + assert found_tools.keys() == expected_tools - # Verify search_nodes tool - search_tool = next(t for t in tools if t.name == "search_nodes") - assert search_tool.inputSchema["required"] == ["query"] - assert "query" in search_tool.inputSchema["properties"] + # Verify schemas include required fields + assert "entities" in found_tools["create_entities"].inputSchema["required"] + assert "query" in found_tools["search_nodes"].inputSchema["required"] @pytest.mark.anyio -async def test_call_create_entities(test_entity_data): - """Test creating an entity through the tool interface.""" - result = await server.call_tool("create_entities", test_entity_data) +async def test_create_entities_camel_case(test_entity_data): + """Test creating an entity with camelCase data (like from MCP).""" + result = await server.handle_call_tool("create_entities", test_entity_data) assert len(result) == 1 assert isinstance(result[0], TextContent) - text_content = result[0].text - assert "Test Entity" in text_content + assert "Test Entity" in result[0].text @pytest.mark.anyio -async def test_call_search_nodes(test_entity_data): +async def test_create_entities_snake_case(test_entity_snake_case): + """Test creating an entity with snake_case data (like internal usage).""" + result = await server.handle_call_tool("create_entities", test_entity_snake_case) + + assert len(result) == 1 + assert isinstance(result[0], TextContent) + assert "Test Entity" in result[0].text + +@pytest.mark.anyio +async def test_search_nodes(test_entity_data): """Test searching for an entity after creating it.""" # First create an entity - await server.call_tool("create_entities", test_entity_data) + await server.handle_call_tool("create_entities", test_entity_data) # Then search for it - result = await server.call_tool("search_nodes", {"query": "Test Entity"}) + result = await server.handle_call_tool("search_nodes", {"query": "Test Entity"}) + + assert len(result) == 1 + assert isinstance(result[0], TextContent) + assert "Test Entity" in result[0].text + +@pytest.mark.anyio +async def test_add_observations(test_entity_data): + """Test adding observations to an existing entity.""" + # First create an entity + await server.handle_call_tool("create_entities", test_entity_data) + + # Add new observations using camelCase + result = await server.handle_call_tool("add_observations", { + "entityId": "Test Entity", + "observations": [{"content": "A new observation"}] + }) assert len(result) == 1 assert isinstance(result[0], TextContent) @@ -62,19 +99,30 @@ async def test_call_search_nodes(test_entity_data): @pytest.mark.anyio async def test_invalid_tool_name(): """Test calling a non-existent tool.""" - with pytest.raises(Exception) as exc: # We could be more specific about error type - await server.call_tool("not_a_tool", {}) + with pytest.raises(Exception) as exc: + await server.handle_call_tool("not_a_tool", {}) assert "Unknown tool" in str(exc.value) @pytest.mark.anyio async def test_invalid_parameters(): - """Test calling tools with invalid parameters.""" - # Test missing required parameter + """Test validation with invalid parameters.""" + # Test missing required field with pytest.raises(Exception) as exc: - await server.call_tool("search_nodes", {}) - assert "query" in str(exc.value) + await server.handle_call_tool("search_nodes", {}) + assert "query" in str(exc.value).lower() # Test empty entities list with pytest.raises(Exception) as exc: - await server.call_tool("create_entities", {"entities": []}) - assert "min_items" in str(exc.value) \ No newline at end of file + await server.handle_call_tool("create_entities", {"entities": []}) + assert "length" in str(exc.value).lower() + + # Test invalid case mixing (should pass due to aliases) + result = await server.handle_call_tool("create_entities", { + "entities": [{ + "name": "Mixed Case Test", + "entityType": "test", # camelCase + "observations": [{"content": "Testing mixed case"}] + }] + }) + assert len(result) == 1 + assert "Mixed Case Test" in result[0].text \ No newline at end of file