tests working for mcp server

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,

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))

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)

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
 
 

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
+    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()

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