From b2615515dfda7cb43297b3f21fb55eec288f1ac1 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sun, 8 Dec 2024 19:32:00 -0600 Subject: [PATCH] tests working for mcp server --- src/basic_memory/deps.py | 6 +- src/basic_memory/mcp/server.py | 195 +++++++------------- src/basic_memory/schemas.py | 47 +++-- src/basic_memory/services/memory_service.py | 4 +- tests/test_mcp_server.py | 113 ++++++++---- tests/test_memory_service.py | 24 ++- 6 files changed, 197 insertions(+), 192 deletions(-) diff --git a/src/basic_memory/deps.py b/src/basic_memory/deps.py index 4339d20e..27b37eb9 100644 --- a/src/basic_memory/deps.py +++ b/src/basic_memory/deps.py @@ -12,15 +12,15 @@ from basic_memory.db import DatabaseType, get_database_url, init_database, get_s async def get_entity_repo(session: AsyncSession) -> EntityRepository: """Get an EntityRepository instance.""" - return EntityRepository(session, DbEntity) + return EntityRepository(session) # Entity type is handled in EntityRepository.__init__ async def get_observation_repo(session: AsyncSession) -> ObservationRepository: """Get an ObservationRepository instance.""" - return ObservationRepository(session, DbObservation) + return ObservationRepository(session) async def get_relation_repo(session: AsyncSession) -> RelationRepository: """Get a RelationRepository instance.""" - return RelationRepository(session, DbRelation) + return RelationRepository(session) async def get_entity_service( project_path: Path, diff --git a/src/basic_memory/mcp/server.py b/src/basic_memory/mcp/server.py index 09b389f6..86ae27a8 100644 --- a/src/basic_memory/mcp/server.py +++ b/src/basic_memory/mcp/server.py @@ -1,19 +1,26 @@ """MCP server implementation for basic-memory.""" -from typing import List, Dict, Any, Optional +from typing import List, Dict, Any, Optional, Literal, Callable, Awaitable +from typing_extensions import TypeAlias from mcp.server import Server from mcp.types import Tool, EmbeddedResource, TextResourceContents, METHOD_NOT_FOUND, INVALID_PARAMS, INTERNAL_ERROR from mcp.shared.exceptions import McpError from pydantic.networks import AnyUrl -from pydantic import TypeAdapter, BaseModel +from pydantic import TypeAdapter, BaseModel, ConfigDict from basic_memory.config import ProjectConfig, create_project_services from basic_memory.fileio import EntityNotFoundError from basic_memory.schemas import ( - ObservationIn, RelationIn, EntityIn, ObservationsIn, + # Tool inputs + CreateEntitiesInput, SearchNodesInput, OpenNodesInput, + AddObservationsInput, CreateRelationsInput, DeleteEntitiesInput, + DeleteObservationsInput, + # Tool responses CreateEntitiesResponse, SearchNodesResponse, OpenNodesResponse, AddObservationsResponse, CreateRelationsResponse, DeleteEntitiesResponse, - DeleteObservationsResponse, EntityOut, ObservationOut + DeleteObservationsResponse, + # Base models + EntityOut, ObservationOut ) from basic_memory.services.memory_service import MemoryService @@ -22,6 +29,19 @@ MIME_TYPE = "application/vnd.basic-memory+json" url_validator = TypeAdapter(AnyUrl) BASIC_MEMORY_URI = url_validator.validate_python("basic-memory://response") +# Define tool name type and handler type +ToolName = Literal[ + "create_entities", + "search_nodes", + "open_nodes", + "add_observations", + "create_relations", + "delete_entities", + "delete_observations" +] + +ToolHandler: TypeAlias = Callable[[MemoryService, Dict[str, Any]], Awaitable[EmbeddedResource]] + def create_response(response: BaseModel) -> EmbeddedResource: """Create standard MCP response from any response model.""" @@ -40,11 +60,11 @@ async def handle_create_entities( 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"]] + # Validate input + input_args = CreateEntitiesInput.model_validate(args) # Call service with validated data - entities = await service.create_entities(entities_data) + entities = await service.create_entities(input_args.entities) # Format response response = CreateEntitiesResponse(entities=[EntityOut.model_validate(entity) for entity in entities]) @@ -56,10 +76,11 @@ async def handle_search_nodes( args: Dict[str, Any] ) -> EmbeddedResource: """Handle search_nodes tool call.""" - results = await service.search_nodes(args["query"]) + input_args = SearchNodesInput.model_validate(args) + results = await service.search_nodes(input_args.query) response = SearchNodesResponse( matches=[EntityOut.model_validate(entity) for entity in results], - query=args["query"] + query=input_args.query ) return create_response(response) @@ -69,7 +90,8 @@ async def handle_open_nodes( args: Dict[str, Any] ) -> EmbeddedResource: """Handle open_nodes tool call.""" - entities = await service.open_nodes(args["names"]) + input_args = OpenNodesInput.model_validate(args) + entities = await service.open_nodes(input_args.names) response = OpenNodesResponse(entities=[EntityOut.model_validate(entity) for entity in entities]) return create_response(response) @@ -80,15 +102,15 @@ async def handle_add_observations( ) -> EmbeddedResource: """Handle add_observations tool call.""" # Validate input - observations_in = ObservationsIn.model_validate(args) + input_args = AddObservationsInput.model_validate(args) # Call service with validated data - observations = await service.add_observations(observations_in) + observations = await service.add_observations(input_args) # Format response response = AddObservationsResponse( - entity_id=observations_in.entity_id, - added_observations=[ObservationOut.model_validate(observation) for observation in observations] + entity_id=input_args.entity_id, + added_observations=[ObservationOut.model_validate(obs) for obs in observations] ) return create_response(response) @@ -98,12 +120,13 @@ async def handle_create_relations( 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"]] + # Validate input + input_args = CreateRelationsInput.model_validate(args) # Call service with validated data - created = await service.create_relations(relations) + created = await service.create_relations(input_args.relations) + # Format response response = CreateRelationsResponse(relations=created) return create_response(response) @@ -113,7 +136,8 @@ async def handle_delete_entities( args: Dict[str, Any] ) -> EmbeddedResource: """Handle delete_entities tool call.""" - deleted = await service.delete_entities(args["names"]) + input_args = DeleteEntitiesInput.model_validate(args) + deleted = await service.delete_entities(input_args.names) response = DeleteEntitiesResponse(deleted=deleted) return create_response(response) @@ -123,7 +147,8 @@ async def handle_delete_observations( args: Dict[str, Any] ) -> EmbeddedResource: """Handle delete_observations tool call.""" - entity, deleted = await service.delete_observations(args["deletions"]) + input_args = DeleteObservationsInput.model_validate(args) + entity, deleted = await service.delete_observations(input_args.deletions) response = DeleteObservationsResponse( entity=entity, deleted=deleted @@ -131,6 +156,18 @@ async def handle_delete_observations( return create_response(response) +# Map tool names to handlers +TOOL_HANDLERS: Dict[ToolName, ToolHandler] = { + "create_entities": handle_create_entities, + "search_nodes": handle_search_nodes, + "open_nodes": handle_open_nodes, + "add_observations": handle_add_observations, + "create_relations": handle_create_relations, + "delete_entities": handle_delete_entities, + "delete_observations": handle_delete_observations, +} + + class MemoryServer(Server): """Extended server class that exposes handlers for testing.""" @@ -149,117 +186,37 @@ class MemoryServer(Server): 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"] - } + inputSchema=CreateEntitiesInput.model_json_schema() ), Tool( name="search_nodes", description="Search for nodes in the knowledge graph", - inputSchema={ - "type": "object", - "properties": { - "query": { - "type": "string", - "minLength": 1 - } - }, - "required": ["query"] - } + inputSchema=SearchNodesInput.model_json_schema() ), 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"] - } + inputSchema=OpenNodesInput.model_json_schema() ), 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"] - } + inputSchema=AddObservationsInput.model_json_schema() ), 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"] - } + inputSchema=CreateRelationsInput.model_json_schema() ), 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"] - } + inputSchema=DeleteEntitiesInput.model_json_schema() ), 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"] - } + inputSchema=DeleteObservationsInput.model_json_schema() ) ] @@ -272,31 +229,17 @@ class MemoryServer(Server): ) -> List[EmbeddedResource]: """Handle tool calls by delegating to the memory service.""" try: + # Check if tool exists + if name not in TOOL_HANDLERS: + raise McpError(METHOD_NOT_FOUND, f"Unknown tool: {name}") + service = await create_project_services( self.config, memory_service=memory_service ) - match name: - case "create_entities": - return [await handle_create_entities(service, arguments)] - case "search_nodes": - return [await handle_search_nodes(service, arguments)] - case "open_nodes": - return [await handle_open_nodes(service, arguments)] - case "add_observations": - return [await handle_add_observations(service, arguments)] - case "create_relations": - return [await handle_create_relations(service, arguments)] - case "delete_entities": - return [await handle_delete_entities(service, arguments)] - case "delete_observations": - return [await handle_delete_observations(service, arguments)] - case _: - raise McpError( - METHOD_NOT_FOUND, - f"Unknown tool: {name}" - ) + tool_name = name # type: ignore + return [await TOOL_HANDLERS[tool_name](service, arguments)] except ValueError as e: raise McpError(INVALID_PARAMS, str(e)) diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index 7de50964..58f3aaae 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -4,11 +4,12 @@ These models define the schema for our core data types while remaining independent from storage/persistence concerns. """ from datetime import datetime, UTC -from typing import List, Optional, Dict, Any +from typing import List, Optional, Dict, Any, Annotated from uuid import uuid4 - +from annotated_types import Gt, Len from pydantic import BaseModel, Field, model_validator, ConfigDict +# Base Models class ObservationIn(BaseModel): """Schema for creating a single observation.""" content: str @@ -40,10 +41,10 @@ class RelationIn(BaseModel): to_id: str = Field(alias="toId") relation_type: str = Field(alias="relationType") context: Optional[str] = None + model_config = ConfigDict(populate_by_name=True) class RelationOut(BaseModel): - """Schema for relation data returned from the service.""" id: int from_id: str = Field(alias="fromId") to_id: str = Field(alias="toId") @@ -52,7 +53,6 @@ class RelationOut(BaseModel): 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") @@ -88,7 +88,38 @@ class EntityOut(EntityBase): relations: List[RelationOut] = [] model_config = ConfigDict(populate_by_name=True, from_attributes=True) -# Tool Response Models +# Tool Input Schemas +class CreateEntitiesInput(BaseModel): + """Input schema for create_entities tool.""" + entities: Annotated[List[EntityIn], Len(min_length=1)] + +class SearchNodesInput(BaseModel): + """Input schema for search_nodes tool.""" + query: str + +class OpenNodesInput(BaseModel): + """Input schema for open_nodes tool.""" + names: Annotated[List[str], Len(min_length=1)] + +class AddObservationsInput(BaseModel): + """Input schema for add_observations tool.""" + entity_id: str = Field(alias="entityId") + observations: List[ObservationIn] + model_config = ConfigDict(populate_by_name=True) + +class CreateRelationsInput(BaseModel): + """Input schema for create_relations tool.""" + relations: List[RelationIn] + +class DeleteEntitiesInput(BaseModel): + """Input schema for delete_entities tool.""" + names: List[str] + +class DeleteObservationsInput(BaseModel): + """Input schema for delete_observations tool.""" + deletions: List[Dict[str, Any]] # TODO: Make this more specific + +# Tool Response Schemas class CreateEntitiesResponse(BaseModel): """Response for create_entities tool.""" entities: List[EntityOut] @@ -125,10 +156,4 @@ 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] 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 56153ede..c7bde557 100644 --- a/src/basic_memory/services/memory_service.py +++ b/src/basic_memory/services/memory_service.py @@ -5,9 +5,9 @@ from pathlib import Path from basic_memory.models import Entity, Observation from basic_memory.schemas import ( - ObservationsIn, ObservationsOut, ObservationOut, EntityIn, EntityOut, RelationIn + ObservationsIn, EntityIn, RelationIn ) -from basic_memory.fileio import write_entity_file, read_entity_file, delete_entity_file +from basic_memory.fileio import write_entity_file, read_entity_file from basic_memory.services import EntityService, RelationService, ObservationService diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index d9abd6a3..3e7c1193 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -1,9 +1,9 @@ """Tests for the MCP server implementation.""" -import json import pytest from pathlib import Path -from mcp.types import EmbeddedResource, TextResourceContents +from mcp.types import EmbeddedResource +from mcp.shared.exceptions import McpError from basic_memory.mcp.server import MemoryServer, MIME_TYPE, BASIC_MEMORY_URI from basic_memory.config import ProjectConfig from basic_memory.schemas import ( @@ -63,8 +63,9 @@ async def test_list_tools(test_config): assert found_tools.keys() == expected_tools # Verify schemas include required fields - assert "entities" in found_tools["create_entities"].inputSchema["required"] - assert "query" in found_tools["search_nodes"].inputSchema["required"] + search_schema = found_tools["search_nodes"].inputSchema + assert "query" in search_schema["properties"] + assert search_schema["required"] == ["query"] @pytest.mark.anyio async def test_create_entities_camel_case(test_entity_data, memory_service, test_config): @@ -79,9 +80,9 @@ async def test_create_entities_camel_case(test_entity_data, memory_service, test assert len(result) == 1 assert isinstance(result[0], EmbeddedResource) assert result[0].type == "resource" - assert isinstance(result[0].resource, TextResourceContents) + assert isinstance(result[0].resource.uri, type(BASIC_MEMORY_URI)) + assert str(result[0].resource.uri) == str(BASIC_MEMORY_URI) 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 @@ -102,9 +103,9 @@ async def test_create_entities_snake_case(test_entity_snake_case, memory_service assert len(result) == 1 assert isinstance(result[0], EmbeddedResource) assert result[0].type == "resource" - assert isinstance(result[0].resource, TextResourceContents) + assert isinstance(result[0].resource.uri, type(BASIC_MEMORY_URI)) + assert str(result[0].resource.uri) == str(BASIC_MEMORY_URI) 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 @@ -134,9 +135,9 @@ async def test_search_nodes(test_entity_data, memory_service, test_config): assert len(result) == 1 assert isinstance(result[0], EmbeddedResource) assert result[0].type == "resource" - assert isinstance(result[0].resource, TextResourceContents) + assert isinstance(result[0].resource.uri, type(BASIC_MEMORY_URI)) + assert str(result[0].resource.uri) == str(BASIC_MEMORY_URI) 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 @@ -171,9 +172,9 @@ async def test_add_observations(test_entity_data, memory_service, test_config): assert len(result) == 1 assert isinstance(result[0], EmbeddedResource) assert result[0].type == "resource" - assert isinstance(result[0].resource, TextResourceContents) + assert isinstance(result[0].resource.uri, type(BASIC_MEMORY_URI)) + assert str(result[0].resource.uri) == str(BASIC_MEMORY_URI) 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 @@ -184,38 +185,70 @@ async def test_add_observations(test_entity_data, memory_service, test_config): async def test_invalid_tool_name(test_config): """Test calling a non-existent tool.""" server_instance = MemoryServer(config=test_config) - with pytest.raises(Exception) as exc: + with pytest.raises(McpError) as exc: await server_instance.handle_call_tool("not_a_tool", {}) assert "Unknown tool" in str(exc.value) @pytest.mark.anyio -async def test_invalid_parameters(test_config): - """Test validation with invalid parameters.""" - server_instance = MemoryServer(config=test_config) +class TestInputValidation: + """Test input validation for various tools.""" - # Test missing required field - with pytest.raises(Exception) as exc: - await server_instance.handle_call_tool("search_nodes", {}) - assert "query" in str(exc.value).lower() + async def test_missing_required_field(self, test_config): + """Test validation when required fields are missing.""" + server_instance = MemoryServer(config=test_config) + + with pytest.raises(McpError) as exc: + await server_instance.handle_call_tool("search_nodes", {}) + assert "query" in str(exc.value).lower() + + with pytest.raises(McpError) as exc: + await server_instance.handle_call_tool("create_entities", {}) + assert "entities" in str(exc.value).lower() - # Test empty entities list - with pytest.raises(Exception) as exc: - await server_instance.handle_call_tool("create_entities", {"entities": []}) - assert "min_items" in str(exc.value).lower() + async def test_empty_arrays(self, test_config): + """Test validation of array fields that can't be empty.""" + server_instance = MemoryServer(config=test_config) + + with pytest.raises(McpError) as exc: + await server_instance.handle_call_tool("create_entities", {"entities": []}) + assert "validation error" in str(exc.value).lower() + + with pytest.raises(McpError) as exc: + await server_instance.handle_call_tool("open_nodes", {"names": []}) + assert "validation error" in str(exc.value).lower() - # Test invalid case mixing (should still work with aliases) - result = await server_instance.handle_call_tool( - "create_entities", - { - "entities": [{ - "name": "Mixed Case Test", - "entityType": "test", # camelCase - "observations": [{"content": "Testing case handling"}] - }] - } - ) - assert len(result) == 1 - 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 + async def test_invalid_field_types(self, test_config): + """Test validation when fields have wrong types.""" + server_instance = MemoryServer(config=test_config) + + with pytest.raises(McpError) as exc: + await server_instance.handle_call_tool("search_nodes", {"query": 123}) + assert "str" in str(exc.value).lower() + + with pytest.raises(McpError) as exc: + await server_instance.handle_call_tool("create_entities", {"entities": "not an array"}) + assert "array" in str(exc.value).lower() or "list" in str(exc.value).lower() + + async def test_invalid_nested_fields(self, test_config): + """Test validation of nested object fields.""" + server_instance = MemoryServer(config=test_config) + + with pytest.raises(McpError) as exc: + await server_instance.handle_call_tool("create_entities", { + "entities": [{ + "name": "Test", + # Missing required entityType + "observations": [] + }] + }) + assert "entitytype" in str(exc.value).lower() + + with pytest.raises(McpError) as exc: + await server_instance.handle_call_tool("add_observations", { + "entityId": "123", + "observations": [{ + # Missing required content field + "context": "test" + }] + }) + assert "content" in str(exc.value).lower() \ No newline at end of file diff --git a/tests/test_memory_service.py b/tests/test_memory_service.py index 90840333..c0da8559 100644 --- a/tests/test_memory_service.py +++ b/tests/test_memory_service.py @@ -3,7 +3,7 @@ import pytest from basic_memory.services import MemoryService from basic_memory.fileio import read_entity_file from basic_memory.models import Entity as EntityModel, Observation, Relation -from basic_memory.schemas import EntityIn +from basic_memory.schemas import EntityIn, CreateEntitiesInput, AddObservationsInput test_entities_data = [ { @@ -21,8 +21,9 @@ test_entities_data = [ @pytest.mark.asyncio async def test_create_entities(memory_service: MemoryService): """Should create multiple entities in parallel with their observations.""" - # Create entities - returns List[models.Entity] - entities = await memory_service.create_entities(test_entities_data) + + entity_input = CreateEntitiesInput.model_validate({"entities": test_entities_data}) + entities = await memory_service.create_entities(entity_input.entities) # Verify the SQLAlchemy models were created assert len(entities) == 2 @@ -51,8 +52,8 @@ async def test_create_entities(memory_service: MemoryService): @pytest.mark.asyncio async def test_add_observations(memory_service: MemoryService): """Should add observations to an existing entity.""" - # First create an entity - returns SQLAlchemy Entity - entities = await memory_service.create_entities([test_entities_data[0]]) + entity_input = CreateEntitiesInput.model_validate({"entities": test_entities_data}) + entities = await memory_service.create_entities([entity_input.entities[0]]) entity = entities[0] # Create observations input @@ -65,7 +66,8 @@ async def test_add_observations(memory_service: MemoryService): } # Add observations - returns List[models.Observation] - added_observations = await memory_service.add_observations(observations_data) + observation_input = AddObservationsInput.model_validate(observations_data) + added_observations = await memory_service.add_observations(observation_input) # Check the SQLAlchemy model results assert len(added_observations) == 2 @@ -94,13 +96,14 @@ async def test_add_observations_nonexistent_entity(memory_service: MemoryService } with pytest.raises(Exception) as exc: # We might want to define a specific error type - await memory_service.add_observations(observations_data) + observation_input = AddObservationsInput.model_validate(observations_data) + await memory_service.add_observations(observation_input) @pytest.mark.asyncio async def test_create_relations(memory_service: MemoryService): """Should create relations between entities and update both filesystem and database.""" - # First create entities - returns List[models.Entity] - entities = await memory_service.create_entities(test_entities_data) + entity_input = CreateEntitiesInput.model_validate({"entities": test_entities_data}) + entities = await memory_service.create_entities(entity_input.entities) entity1, entity2 = entities # Create test relations data using actual entity IDs @@ -181,7 +184,8 @@ async def test_create_relations(memory_service: MemoryService): async def test_create_relations_with_invalid_entity_id(memory_service: MemoryService): """Should raise an appropriate error when trying to create relations with non-existent entity IDs.""" # Create one entity - returns SQLAlchemy Entity - entities = await memory_service.create_entities([test_entities_data[0]]) + entity_input = CreateEntitiesInput.model_validate({"entities": test_entities_data}) + entities = await memory_service.create_entities([entity_input.entities[0]]) entity1 = entities[0] # Try to create relation with non-existent entity ID