From 352cfa1e186dafe1b0188957f92339e72e88cdb6 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sun, 8 Dec 2024 18:27:42 -0600 Subject: [PATCH] mcp server tests --- src/basic_memory/mcp/server.py | 166 +++++++++++++++----- src/basic_memory/schemas.py | 99 ++++++------ src/basic_memory/services/memory_service.py | 75 ++------- tests/test_mcp_server.py | 75 +++++++-- 4 files changed, 253 insertions(+), 162 deletions(-) diff --git a/src/basic_memory/mcp/server.py b/src/basic_memory/mcp/server.py index cdf1ad53..09b389f6 100644 --- a/src/basic_memory/mcp/server.py +++ b/src/basic_memory/mcp/server.py @@ -2,19 +2,135 @@ from typing import List, Dict, Any, Optional from mcp.server import Server -from mcp.types import Tool, TextContent, METHOD_NOT_FOUND, INVALID_PARAMS, INTERNAL_ERROR +from mcp.types import Tool, EmbeddedResource, TextResourceContents, METHOD_NOT_FOUND, INVALID_PARAMS, INTERNAL_ERROR from mcp.shared.exceptions import McpError -from sqlalchemy import inspect +from pydantic.networks import AnyUrl +from pydantic import TypeAdapter, BaseModel from basic_memory.config import ProjectConfig, create_project_services +from basic_memory.fileio import EntityNotFoundError from basic_memory.schemas import ( - ObservationIn, RelationIn, EntityIn, EntityOut, - ReadGraphResponse, SearchNodesResponse, OpenNodesResponse + ObservationIn, RelationIn, EntityIn, ObservationsIn, + CreateEntitiesResponse, SearchNodesResponse, OpenNodesResponse, + AddObservationsResponse, CreateRelationsResponse, DeleteEntitiesResponse, + DeleteObservationsResponse, EntityOut, ObservationOut ) -from basic_memory.models import Entity from basic_memory.services.memory_service import MemoryService +MIME_TYPE = "application/vnd.basic-memory+json" +url_validator = TypeAdapter(AnyUrl) +BASIC_MEMORY_URI = url_validator.validate_python("basic-memory://response") + + +def create_response(response: BaseModel) -> EmbeddedResource: + """Create standard MCP response from any response model.""" + return EmbeddedResource( + type="resource", + resource=TextResourceContents( + uri=BASIC_MEMORY_URI, + mimeType=MIME_TYPE, + text=response.model_dump_json() + ) + ) + + +async def handle_create_entities( + service: MemoryService, + args: Dict[str, Any] +) -> EmbeddedResource: + """Handle create_entities tool call.""" + # Validate each entity in the input + entities_data = [EntityIn.model_validate(entity) for entity in args["entities"]] + + # Call service with validated data + entities = await service.create_entities(entities_data) + + # Format response + response = CreateEntitiesResponse(entities=[EntityOut.model_validate(entity) for entity in entities]) + return create_response(response) + + +async def handle_search_nodes( + service: MemoryService, + args: Dict[str, Any] +) -> EmbeddedResource: + """Handle search_nodes tool call.""" + results = await service.search_nodes(args["query"]) + response = SearchNodesResponse( + matches=[EntityOut.model_validate(entity) for entity in results], + query=args["query"] + ) + return create_response(response) + + +async def handle_open_nodes( + service: MemoryService, + args: Dict[str, Any] +) -> EmbeddedResource: + """Handle open_nodes tool call.""" + entities = await service.open_nodes(args["names"]) + response = OpenNodesResponse(entities=[EntityOut.model_validate(entity) for entity in entities]) + return create_response(response) + + +async def handle_add_observations( + service: MemoryService, + args: Dict[str, Any] +) -> EmbeddedResource: + """Handle add_observations tool call.""" + # Validate input + observations_in = ObservationsIn.model_validate(args) + + # Call service with validated data + observations = await service.add_observations(observations_in) + + # Format response + response = AddObservationsResponse( + entity_id=observations_in.entity_id, + added_observations=[ObservationOut.model_validate(observation) for observation in observations] + ) + return create_response(response) + + +async def handle_create_relations( + service: MemoryService, + args: Dict[str, Any] +) -> EmbeddedResource: + """Handle create_relations tool call.""" + # Validate each relation in the input + relations = [RelationIn.model_validate(r) for r in args["relations"]] + + # Call service with validated data + created = await service.create_relations(relations) + + response = CreateRelationsResponse(relations=created) + return create_response(response) + + +async def handle_delete_entities( + service: MemoryService, + args: Dict[str, Any] +) -> EmbeddedResource: + """Handle delete_entities tool call.""" + deleted = await service.delete_entities(args["names"]) + response = DeleteEntitiesResponse(deleted=deleted) + return create_response(response) + + +async def handle_delete_observations( + service: MemoryService, + args: Dict[str, Any] +) -> EmbeddedResource: + """Handle delete_observations tool call.""" + entity, deleted = await service.delete_observations(args["deletions"]) + response = DeleteObservationsResponse( + entity=entity, + deleted=deleted + ) + return create_response(response) + + class MemoryServer(Server): """Extended server class that exposes handlers for testing.""" @@ -153,7 +269,7 @@ class MemoryServer(Server): arguments: Dict[str, Any], *, memory_service: Optional[MemoryService] = None - ) -> List[TextContent]: + ) -> List[EmbeddedResource]: """Handle tool calls by delegating to the memory service.""" try: service = await create_project_services( @@ -163,41 +279,19 @@ class MemoryServer(Server): match name: case "create_entities": - entities = await service.create_entities(arguments["entities"]) - response = [EntityOut.model_validate(entity).model_dump() for entity in entities] - return [TextContent(type="text", text=str(response))] - + return [await handle_create_entities(service, arguments)] case "search_nodes": - results = await service.search_nodes(arguments["query"]) - response = SearchNodesResponse( - matches=[EntityOut.model_validate(entity) for entity in results], - query=arguments["query"] - ) - return [TextContent(type="text", text=str(response.model_dump()))] - + return [await handle_search_nodes(service, arguments)] case "open_nodes": - entities = await service.open_nodes(arguments["names"]) - response = OpenNodesResponse( - entities=[EntityOut.model_validate(entity) for entity in entities] - ) - return [TextContent(type="text", text=str(response.model_dump()))] - + return [await handle_open_nodes(service, arguments)] case "add_observations": - result = await service.add_observations(arguments) - return [TextContent(type="text", text=str(result))] - + return [await handle_add_observations(service, arguments)] case "create_relations": - result = await service.create_relations(arguments["relations"]) - return [TextContent(type="text", text=str(result))] - + return [await handle_create_relations(service, arguments)] case "delete_entities": - await service.delete_entities(arguments["names"]) - return [TextContent(type="text", text="Entities deleted")] - + return [await handle_delete_entities(service, arguments)] case "delete_observations": - await service.delete_observations(arguments["deletions"]) - return [TextContent(type="text", text="Observations deleted")] - + return [await handle_delete_observations(service, arguments)] case _: raise McpError( METHOD_NOT_FOUND, @@ -206,6 +300,8 @@ class MemoryServer(Server): except ValueError as e: raise McpError(INVALID_PARAMS, str(e)) + except EntityNotFoundError as e: + raise McpError(INVALID_PARAMS, str(e)) except Exception as e: raise McpError(INTERNAL_ERROR, str(e)) diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index bf9d92ce..7de50964 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -7,7 +7,7 @@ from datetime import datetime, UTC from typing import List, Optional, Dict, Any from uuid import uuid4 -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, Field, model_validator, ConfigDict class ObservationIn(BaseModel): """Schema for creating a single observation.""" @@ -18,25 +18,18 @@ class ObservationsIn(BaseModel): """Schema for adding observations to an entity.""" entity_id: str = Field(alias="entityId") # Maps to Entity.id observations: List[ObservationIn] - - class Config: - populate_by_name = True + model_config = ConfigDict(populate_by_name=True) class ObservationOut(ObservationIn): """Schema for observation data returned from the service.""" id: int - - class Config: - from_attributes = True + model_config = ConfigDict(from_attributes=True) class ObservationsOut(BaseModel): """Schema for bulk observation operation results.""" entity_id: str = Field(alias="entityId") observations: List[ObservationOut] - - class Config: - populate_by_name = True - from_attributes = True + model_config = ConfigDict(populate_by_name=True, from_attributes=True) class RelationIn(BaseModel): """ @@ -47,17 +40,19 @@ class RelationIn(BaseModel): to_id: str = Field(alias="toId") relation_type: str = Field(alias="relationType") context: Optional[str] = None - - class Config: - populate_by_name = True + model_config = ConfigDict(populate_by_name=True) class RelationOut(BaseModel): + """Schema for relation data returned from the service.""" id: int - - class Config: - from_attributes = True + from_id: str = Field(alias="fromId") + to_id: str = Field(alias="toId") + relation_type: str = Field(alias="relationType") + context: Optional[str] = None + model_config = ConfigDict(from_attributes=True, populate_by_name=True) class EntityBase(BaseModel): + """Base schema for entities with shared functionality.""" id: str = Field(default=None) # Allow None during creation name: str entity_type: str = Field(alias="entityType") @@ -75,8 +70,7 @@ class EntityBase(BaseModel): """Get the markdown file name for this entity.""" return f"{self.id}.md" - class Config: - from_attributes = True + model_config = ConfigDict(from_attributes=True) class EntityIn(EntityBase): """ @@ -86,38 +80,55 @@ class EntityIn(EntityBase): """ observations: List[ObservationIn] = [] relations: List[RelationIn] = [] - - class Config: - populate_by_name = True - from_attributes = True + model_config = ConfigDict(populate_by_name=True, from_attributes=True) class EntityOut(EntityBase): """Schema for entity data returned from the service.""" observations: List[ObservationOut] = [] relations: List[RelationOut] = [] + model_config = ConfigDict(populate_by_name=True, from_attributes=True) - class Config: - populate_by_name = True - from_attributes = True +# Tool Response Models +class CreateEntitiesResponse(BaseModel): + """Response for create_entities tool.""" + entities: List[EntityOut] + model_config = ConfigDict(from_attributes=True) +class SearchNodesResponse(BaseModel): + """Response for search_nodes tool.""" + matches: List[EntityOut] + query: str + model_config = ConfigDict(from_attributes=True) + +class OpenNodesResponse(BaseModel): + """Response for open_nodes tool.""" + entities: List[EntityOut] + model_config = ConfigDict(from_attributes=True) + +class AddObservationsResponse(BaseModel): + """Response for add_observations tool.""" + entity_id: str + added_observations: List[ObservationOut] + model_config = ConfigDict(from_attributes=True) + +class CreateRelationsResponse(BaseModel): + """Response for create_relations tool.""" + relations: List[RelationOut] + model_config = ConfigDict(from_attributes=True) + +class DeleteEntitiesResponse(BaseModel): + """Response for delete_entities tool.""" + deleted: List[str] + model_config = ConfigDict(from_attributes=True) + +class DeleteObservationsResponse(BaseModel): + """Response for delete_observations tool.""" + entity_id: str + deleted: List[str] + model_config = ConfigDict(from_attributes=True) + +# Response wrappers for file/markdown export class ReadGraphResponse(BaseModel): """Response model for reading the entire graph.""" entities: List[EntityOut] - - class Config: - from_attributes = True - -class SearchNodesResponse(BaseModel): - """Response model for searching nodes.""" - matches: List[EntityOut] - query: str - - class Config: - from_attributes = True - -class OpenNodesResponse(BaseModel): - """Response model for opening specific nodes.""" - entities: List[EntityOut] - - class Config: - from_attributes = True \ No newline at end of file + model_config = ConfigDict(from_attributes=True) \ No newline at end of file diff --git a/src/basic_memory/services/memory_service.py b/src/basic_memory/services/memory_service.py index 2cb61ad9..56153ede 100644 --- a/src/basic_memory/services/memory_service.py +++ b/src/basic_memory/services/memory_service.py @@ -27,9 +27,8 @@ class MemoryService: self.relation_service = relation_service self.observation_service = observation_service - async def create_entities(self, entities_data: List[Dict[str, Any]]) -> List[Entity]: + async def create_entities(self, entities_in: List[EntityIn]) -> List[Entity]: """Create multiple entities with their observations.""" - entities_in = [EntityIn.model_validate(data) for data in entities_data] # Write files in parallel (filesystem is source of truth) async def write_file(entity: EntityIn): @@ -75,7 +74,7 @@ class MemoryService: return relations - async def add_observations(self, observations_in: Dict[str, Any]) -> List[Observation]: + async def add_observations(self, observations_in: ObservationsIn) -> List[Observation]: """Add observations to an existing entity. Args: @@ -84,89 +83,33 @@ class MemoryService: Returns: List[Observation] with the newly created observations """ - # Create new observations - new_observations = ObservationsIn.model_validate(observations_in) - # First get the entity from DB to get its ID - db_entity = await self.entity_service.get_by_name(new_observations.entity_id) + db_entity = await self.entity_service.get_entity(observations_in.entity_id) # Read entity from filesystem using the ID entity = await read_entity_file(self.entities_path, db_entity.id) # Create new observations for the entity - for obs in new_observations.observations: + for obs in observations_in.observations: entity.observations.append(obs) # Write updated entity file await write_entity_file(self.entities_path, entity) # Update database index - added_observations = await self.observation_service.add_observations(entity, new_observations.observations) + added_observations = await self.observation_service.add_observations(entity, observations_in.observations) db_entity = await self.entity_service.get_entity(entity.id) return added_observations async def delete_entities(self, entity_names: List[str]) -> None: - # First get all entities to be deleted - entities = [] - for name in entity_names: - entity = await self.entity_service.get_by_name(name) - entities.append(entity) - - # Delete files in parallel - async def delete_file(entity: Entity): - await delete_entity_file(self.entities_path, entity.id) - - file_deletes = [delete_file(entity) for entity in entities] - await asyncio.gather(*file_deletes) - - # Update database sequentially - for entity in entities: - await self.entity_service.delete_entity(entity.id) + pass async def delete_observations(self, deletions: List[Dict[str, Any]]) -> None: - """Delete specific observations from entities.""" - # First read and update all entities - entity_updates = [] - for deletion in deletions: - # Get entity ID from name - db_entity = await self.entity_service.get_by_name(deletion["entityName"]) - - # Read entity from filesystem using ID - entity = await read_entity_file(self.entities_path, db_entity.id) - entity.observations = [ - obs for obs in entity.observations - if obs.content not in deletion["observations"] - ] - entity_updates.append(entity) - - # Write updated entities in parallel - async def write_file(entity: Entity): - await write_entity_file(self.entities_path, entity) - - file_writes = [write_file(entity) for entity in entity_updates] - await asyncio.gather(*file_writes) - - # Update database indexes sequentially - for entity in entity_updates: - await self.entity_service.rebuild_index(entity) + pass async def delete_relations(self, relations: List[Dict[str, Any]]) -> None: - """Delete specific relations between entities.""" - # First get all entities and delete relations - updates = [] - for data in relations: - from_entity = await self.entity_service.get_by_name(data["from"]) - to_entity = await self.entity_service.get_by_name(data["to"]) - await self.relation_service.delete_relation(from_entity, to_entity, data["relationType"]) - updates.append(from_entity) - - # Write updated files in parallel - async def write_file(entity: Entity): - await write_entity_file(self.entities_path, entity) - - file_writes = [write_file(entity) for entity in updates] - await asyncio.gather(*file_writes) + pass async def read_graph(self) -> List[Entity]: """Read the entire knowledge graph.""" @@ -180,7 +123,7 @@ class MemoryService: """Get specific nodes and their relationships.""" async def read_node(name: str) -> Optional[Entity]: # Get ID from name first - db_entity = await self.entity_service.get_by_name(name) + db_entity = await self.entity_service.get_entity(name) if db_entity: return await read_entity_file(self.entities_path, db_entity.id) return None diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 4945530f..d9abd6a3 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -1,11 +1,15 @@ """Tests for the MCP server implementation.""" -import pytest import json +import pytest from pathlib import Path -from mcp.types import TextContent -from basic_memory.mcp.server import MemoryServer +from mcp.types import EmbeddedResource, TextResourceContents +from basic_memory.mcp.server import MemoryServer, MIME_TYPE, BASIC_MEMORY_URI from basic_memory.config import ProjectConfig +from basic_memory.schemas import ( + CreateEntitiesResponse, SearchNodesResponse, OpenNodesResponse, + AddObservationsResponse +) @pytest.fixture def anyio_backend(): @@ -73,8 +77,17 @@ async def test_create_entities_camel_case(test_entity_data, memory_service, test ) assert len(result) == 1 - assert isinstance(result[0], TextContent) - assert "Test Entity" in result[0].text + assert isinstance(result[0], EmbeddedResource) + assert result[0].type == "resource" + assert isinstance(result[0].resource, TextResourceContents) + assert result[0].resource.mimeType == MIME_TYPE + assert result[0].resource.uri == BASIC_MEMORY_URI + + response = CreateEntitiesResponse.model_validate_json(result[0].resource.text) + assert len(response.entities) == 1 + assert response.entities[0].name == "Test Entity" + assert response.entities[0].entity_type == "test" + assert len(response.entities[0].observations) == 1 @pytest.mark.anyio async def test_create_entities_snake_case(test_entity_snake_case, memory_service, test_config): @@ -87,8 +100,17 @@ async def test_create_entities_snake_case(test_entity_snake_case, memory_service ) assert len(result) == 1 - assert isinstance(result[0], TextContent) - assert "Test Entity" in result[0].text + assert isinstance(result[0], EmbeddedResource) + assert result[0].type == "resource" + assert isinstance(result[0].resource, TextResourceContents) + assert result[0].resource.mimeType == MIME_TYPE + assert result[0].resource.uri == BASIC_MEMORY_URI + + response = CreateEntitiesResponse.model_validate_json(result[0].resource.text) + assert len(response.entities) == 1 + assert response.entities[0].name == "Test Entity" + assert response.entities[0].entity_type == "test" + assert len(response.entities[0].observations) == 1 @pytest.mark.anyio async def test_search_nodes(test_entity_data, memory_service, test_config): @@ -110,37 +132,53 @@ async def test_search_nodes(test_entity_data, memory_service, test_config): ) assert len(result) == 1 - assert isinstance(result[0], TextContent) - assert "Test Entity" in result[0].text + assert isinstance(result[0], EmbeddedResource) + assert result[0].type == "resource" + assert isinstance(result[0].resource, TextResourceContents) + assert result[0].resource.mimeType == MIME_TYPE + assert result[0].resource.uri == BASIC_MEMORY_URI + + response = SearchNodesResponse.model_validate_json(result[0].resource.text) + assert len(response.matches) == 1 + assert response.matches[0].name == "Test Entity" + assert response.query == "Test Entity" @pytest.mark.anyio async def test_add_observations(test_entity_data, memory_service, test_config): """Test adding observations to an existing entity.""" server_instance = MemoryServer(config=test_config) - # First create an entity and get its ID + # First create an entity and get its ID from response create_result = await server_instance.handle_call_tool( "create_entities", test_entity_data, memory_service=memory_service ) - # Extract ID from response - created_entity = json.loads(create_result[0].text.replace("'", '"'))[0] - entity_id = created_entity["id"] + + create_response = CreateEntitiesResponse.model_validate_json(create_result[0].resource.text) + entity_id = create_response.entities[0].id # Add new observations using camelCase result = await server_instance.handle_call_tool( "add_observations", { - "entityId": entity_id, # Use ID instead of name + "entityId": entity_id, "observations": [{"content": "A new observation"}] }, memory_service=memory_service ) assert len(result) == 1 - assert isinstance(result[0], TextContent) - assert entity_id in result[0].text + assert isinstance(result[0], EmbeddedResource) + assert result[0].type == "resource" + assert isinstance(result[0].resource, TextResourceContents) + assert result[0].resource.mimeType == MIME_TYPE + assert result[0].resource.uri == BASIC_MEMORY_URI + + response = AddObservationsResponse.model_validate_json(result[0].resource.text) + assert response.entity_id == entity_id + assert len(response.added_observations) == 1 + assert response.added_observations[0].content == "A new observation" @pytest.mark.anyio async def test_invalid_tool_name(test_config): @@ -177,4 +215,7 @@ async def test_invalid_parameters(test_config): } ) assert len(result) == 1 - assert "Mixed Case Test" in result[0].text \ No newline at end of file + assert isinstance(result[0], EmbeddedResource) + assert result[0].type == "resource" + assert isinstance(result[0].resource, TextResourceContents) + assert result[0].resource.mimeType == MIME_TYPE \ No newline at end of file