From 2b4bd19da1c570ee43b2bcd0b73edeae11eb0faa Mon Sep 17 00:00:00 2001 From: phernandez Date: Sat, 14 Dec 2024 19:09:56 -0600 Subject: [PATCH] cleaned up schemas for api/mcp server --- src/basic_memory/api/routers/knowledge.py | 30 +++++++-------- src/basic_memory/fileio.py | 4 +- src/basic_memory/mcp/server.py | 38 +++++++++---------- src/basic_memory/schemas.py | 31 +++++++-------- src/basic_memory/services/memory_service.py | 8 ++-- src/basic_memory/services/relation_service.py | 5 +-- tests/test_mcp_server.py | 4 +- tests/test_memory_service.py | 18 ++++----- tests/test_relation_service.py | 6 +-- tests/test_schemas.py | 26 ++++++------- 10 files changed, 81 insertions(+), 89 deletions(-) diff --git a/src/basic_memory/api/routers/knowledge.py b/src/basic_memory/api/routers/knowledge.py index e81ff45b..a2206983 100644 --- a/src/basic_memory/api/routers/knowledge.py +++ b/src/basic_memory/api/routers/knowledge.py @@ -4,13 +4,13 @@ from fastapi import APIRouter from basic_memory.deps import MemoryServiceDep from basic_memory.schemas import ( - CreateEntitiesInput, CreateEntitiesResponse, - SearchNodesInput, SearchNodesResponse, - CreateRelationsInput, CreateRelationsResponse, - EntityOut, RelationOut, ObservationsIn, ObservationsOut, ObservationOut, - OpenNodesInput, OpenNodesResponse, + CreateEntitiesRequest, CreateEntitiesResponse, + SearchNodesRequest, SearchNodesResponse, + CreateRelationsRequest, CreateRelationsResponse, + EntityOut, RelationOut, AddObservationsRequest, ObservationOut, + OpenNodesRequest, OpenNodesResponse, DeleteEntitiesResponse, - DeleteObservationsInput, DeleteObservationsResponse + DeleteObservationsRequest, DeleteObservationsResponse, AddObservationsResponse ) router = APIRouter(prefix="/knowledge", tags=["knowledge"]) @@ -18,7 +18,7 @@ router = APIRouter(prefix="/knowledge", tags=["knowledge"]) @router.post("/entities", response_model=CreateEntitiesResponse) async def create_entities( - data: CreateEntitiesInput, + data: CreateEntitiesRequest, memory_service: MemoryServiceDep ) -> CreateEntitiesResponse: """Create new entities in the knowledge graph.""" @@ -48,7 +48,7 @@ async def delete_entity( @router.post("/nodes", response_model=OpenNodesResponse) async def open_nodes( - data: OpenNodesInput, + data: OpenNodesRequest, memory_service: MemoryServiceDep ) -> OpenNodesResponse: """Open specific nodes by their names.""" @@ -58,7 +58,7 @@ async def open_nodes( @router.post("/relations", response_model=CreateRelationsResponse) async def create_relations( - data: CreateRelationsInput, + data: CreateRelationsRequest, memory_service: MemoryServiceDep ) -> CreateRelationsResponse: """Create relations between entities.""" @@ -76,19 +76,19 @@ async def delete_relation( raise NotImplementedError("Delete relation not implemented yet") -@router.post("/observations", response_model=ObservationsOut) +@router.post("/observations", response_model=AddObservationsResponse) async def add_observations( - data: ObservationsIn, + data: AddObservationsRequest, memory_service: MemoryServiceDep -) -> ObservationsOut: +) -> AddObservationsResponse: """Add observations to an entity.""" observations = await memory_service.add_observations(data) - return ObservationsOut(entity_id=data.entity_id, observations=[ObservationOut.model_validate(observation) for observation in observations]) + return AddObservationsResponse(entity_id=data.entity_id, observations=[ObservationOut.model_validate(observation) for observation in observations]) @router.delete("/observations", response_model=DeleteObservationsResponse) async def delete_observations( - data: DeleteObservationsInput, + data: DeleteObservationsRequest, memory_service: MemoryServiceDep ) -> DeleteObservationsResponse: """Delete observations from an entity.""" @@ -98,7 +98,7 @@ async def delete_observations( @router.post("/search", response_model=SearchNodesResponse) async def search_nodes( - data: SearchNodesInput, + data: SearchNodesRequest, memory_service: MemoryServiceDep ) -> SearchNodesResponse: """Search for entities in the knowledge graph.""" diff --git a/src/basic_memory/fileio.py b/src/basic_memory/fileio.py index abbb3661..b00699c1 100644 --- a/src/basic_memory/fileio.py +++ b/src/basic_memory/fileio.py @@ -6,7 +6,7 @@ from pathlib import Path from loguru import logger -from basic_memory.schemas import EntityIn, Relation +from basic_memory.schemas import EntityIn, RelationIn class FileOperationError(Exception): @@ -174,7 +174,7 @@ async def read_entity_file(project_entities_path: Path, entity_id: str) -> Entit relation_type = parts[0] context = parts[1] if len(parts) > 1 else None - relations.append(Relation( # pyright: ignore [reportCallIssue] + relations.append(RelationIn( # pyright: ignore [reportCallIssue] from_id=entity_id, # pyright: ignore [reportCallIssue] to_id=target_id, # pyright: ignore [reportCallIssue] relation_type=relation_type, # pyright: ignore [reportCallIssue] diff --git a/src/basic_memory/mcp/server.py b/src/basic_memory/mcp/server.py index 3c13071b..25a41c6f 100644 --- a/src/basic_memory/mcp/server.py +++ b/src/basic_memory/mcp/server.py @@ -21,15 +21,13 @@ from basic_memory.repository.observation_repository import ObservationRepository from basic_memory.repository.relation_repository import RelationRepository from basic_memory.schemas import ( # Tool inputs - CreateEntitiesInput, SearchNodesInput, OpenNodesInput, - AddObservationsInput, CreateRelationsInput, DeleteEntitiesInput, - DeleteObservationsInput, + CreateEntitiesRequest, SearchNodesRequest, OpenNodesRequest, + CreateRelationsRequest, DeleteEntitiesRequest, + DeleteObservationsRequest, # Tool responses CreateEntitiesResponse, SearchNodesResponse, OpenNodesResponse, AddObservationsResponse, CreateRelationsResponse, DeleteEntitiesResponse, - DeleteObservationsResponse, - # Base models - EntityOut, ObservationOut, RelationOut, ObservationsIn + EntityOut, ObservationOut, RelationOut, AddObservationsRequest ) from basic_memory.services import EntityService, ObservationService, RelationService from basic_memory.services.memory_service import MemoryService @@ -105,7 +103,7 @@ async def handle_create_entities( """Handle create_entities tool call.""" # Validate input logger.debug(f"Creating entities with args: {args}") - input_args = CreateEntitiesInput.model_validate(args) + input_args = CreateEntitiesRequest.model_validate(args) logger.debug(f"Validated input: {len(input_args.entities)} entities") # Call service with validated data @@ -124,7 +122,7 @@ async def handle_search_nodes( ) -> EmbeddedResource: """Handle search_nodes tool call.""" logger.debug(f"Searching nodes with query: {args.get('query')}") - input_args = SearchNodesInput.model_validate(args) + input_args = SearchNodesRequest.model_validate(args) results = await service.search_nodes(input_args.query) logger.debug(f"Found {len(results)} matches for query '{input_args.query}'") response = SearchNodesResponse( @@ -140,7 +138,7 @@ async def handle_open_nodes( ) -> EmbeddedResource: """Handle open_nodes tool call.""" logger.debug(f"Opening nodes: {args.get('names')}") - input_args = OpenNodesInput.model_validate(args) + input_args = OpenNodesRequest.model_validate(args) entities = await service.open_nodes(input_args.names) logger.debug(f"Opened {len(entities)} entities") response = OpenNodesResponse(entities=[EntityOut.model_validate(entity) for entity in entities]) @@ -154,7 +152,7 @@ async def handle_add_observations( """Handle add_observations tool call.""" # Validate input logger.debug(f"Adding observations: {args}") - input_args = ObservationsIn.model_validate(args) + input_args = AddObservationsRequest.model_validate(args) logger.debug(f"Adding {len(input_args.observations)} observations to entity {input_args.entity_id}") # Call service with validated data @@ -164,7 +162,7 @@ async def handle_add_observations( # Format response response = AddObservationsResponse( entity_id=input_args.entity_id, - added_observations=[ObservationOut.model_validate(obs) for obs in observations] + observations=[ObservationOut.model_validate(obs) for obs in observations] ) return create_response(response) @@ -176,7 +174,7 @@ async def handle_create_relations( """Handle create_relations tool call.""" # Validate input logger.debug(f"Creating relations: {args}") - input_args = CreateRelationsInput.model_validate(args) + input_args = CreateRelationsRequest.model_validate(args) logger.debug(f"Creating {len(input_args.relations)} relations") # Call service with validated data @@ -194,7 +192,7 @@ async def handle_delete_entities( ) -> EmbeddedResource: """Handle delete_entities tool call.""" logger.debug(f"Deleting entities: {args}") - input_args = DeleteEntitiesInput.model_validate(args) + input_args = DeleteEntitiesRequest.model_validate(args) deleted = await service.delete_entities(input_args.names) logger.debug(f"Deleted entities: {deleted}") response = DeleteEntitiesResponse(deleted=deleted) @@ -242,37 +240,37 @@ class MemoryServer(Server): Tool( name="create_entities", description="Create multiple new entities in the knowledge graph", - inputSchema=CreateEntitiesInput.model_json_schema() + inputSchema=CreateEntitiesRequest.model_json_schema() ), Tool( name="search_nodes", description="Search for nodes in the knowledge graph", - inputSchema=SearchNodesInput.model_json_schema() + inputSchema=SearchNodesRequest.model_json_schema() ), Tool( name="open_nodes", description="Open specific nodes by their names", - inputSchema=OpenNodesInput.model_json_schema() + inputSchema=OpenNodesRequest.model_json_schema() ), Tool( name="add_observations", description="Add observations to existing entities", - inputSchema=AddObservationsInput.model_json_schema() + inputSchema=AddObservationsRequest.model_json_schema() ), Tool( name="create_relations", description="Create relations between entities", - inputSchema=CreateRelationsInput.model_json_schema() + inputSchema=CreateRelationsRequest.model_json_schema() ), Tool( name="delete_entities", description="Delete entities from the knowledge graph", - inputSchema=DeleteEntitiesInput.model_json_schema() + inputSchema=DeleteEntitiesRequest.model_json_schema() ), Tool( name="delete_observations", description="Delete observations from entities", - inputSchema=DeleteObservationsInput.model_json_schema() + inputSchema=DeleteObservationsRequest.model_json_schema() ) ] logger.debug(f"Returning {len(tools)} available tools") diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index 2febbb39..2a793fd3 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -1,5 +1,5 @@ """Core pydantic models for basic-memory entities, observations, and relations.""" -from typing import List, Optional, Dict, Any, Annotated +from typing import List, Optional, Annotated from annotated_types import Len from pydantic import BaseModel, ConfigDict @@ -9,7 +9,7 @@ class SQLAlchemyOut(BaseModel): model_config = ConfigDict(from_attributes=True) # Base Models -class ObservationsIn(BaseModel): +class AddObservationsRequest(BaseModel): """Schema for adding observations to an entity.""" entity_id: str context: Optional[str] = None @@ -27,7 +27,7 @@ class ObservationsOut(SQLAlchemyOut): observations: List[ObservationOut] model_config = ConfigDict(populate_by_name=True) -class Relation(BaseModel): +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.) @@ -67,7 +67,7 @@ class EntityIn(EntityBase): associated observations. """ observations: List[str] = [] - relations: List[Relation] = [] + relations: List[RelationIn] = [] model_config = ConfigDict(populate_by_name=True) class EntityOut(EntityBase, SQLAlchemyOut): @@ -77,36 +77,31 @@ class EntityOut(EntityBase, SQLAlchemyOut): model_config = ConfigDict(populate_by_name=True) # Tool Input Schemas -class CreateEntitiesInput(BaseModel): +class CreateEntitiesRequest(BaseModel): """Input schema for create_entities tool.""" entities: Annotated[List[EntityIn], Len(min_length=1)] -class SearchNodesInput(BaseModel): +class SearchNodesRequest(BaseModel): """Input schema for search_nodes tool.""" query: str -class OpenNodesInput(BaseModel): +class OpenNodesRequest(BaseModel): """Input schema for open_nodes tool.""" names: Annotated[List[str], Len(min_length=1)] -class AddObservationsInput(ObservationsIn): - """Input schema for add_observations tool.""" - pass - -class CreateRelationsInput(BaseModel): +class CreateRelationsRequest(BaseModel): """Input schema for create_relations tool.""" - relations: List[Relation] + relations: List[RelationIn] -class DeleteEntitiesInput(BaseModel): +class DeleteEntitiesRequest(BaseModel): """Input schema for delete_entities tool.""" names: List[str] -class DeleteObservationsInput(BaseModel): +class DeleteObservationsRequest(BaseModel): """Input schema for delete_observations tool.""" entity_id: str - deletions: List[Dict[str, Any]] # TODO: Make this more specific + deletions: List[str] # TODO: Make this more specific -# Tool Response Schemas class CreateEntitiesResponse(SQLAlchemyOut): """Response for create_entities tool.""" entities: List[EntityOut] @@ -123,7 +118,7 @@ class OpenNodesResponse(SQLAlchemyOut): class AddObservationsResponse(SQLAlchemyOut): """Response for add_observations tool.""" entity_id: str - added_observations: List[ObservationOut] + observations: List[ObservationOut] class CreateRelationsResponse(SQLAlchemyOut): """Response for create_relations tool.""" diff --git a/src/basic_memory/services/memory_service.py b/src/basic_memory/services/memory_service.py index 948b2974..a3d2e25b 100644 --- a/src/basic_memory/services/memory_service.py +++ b/src/basic_memory/services/memory_service.py @@ -3,9 +3,9 @@ import asyncio from typing import List, Dict, Any, Optional, Sequence from pathlib import Path -from basic_memory.models import Entity, Observation, Relation +from basic_memory.models import Entity, Observation from basic_memory.schemas import ( - ObservationsIn, EntityIn, Relation + AddObservationsRequest, EntityIn, RelationIn ) from basic_memory.fileio import write_entity_file, read_entity_file, EntityNotFoundError from basic_memory.services import EntityService, RelationService, ObservationService @@ -112,7 +112,7 @@ class MemoryService: logger.debug(f"Found entity {entity}") return entity - async def create_relations(self, relations_data: List[Relation]) -> List[Relation]: + async def create_relations(self, relations_data: List[RelationIn]) -> List[RelationIn]: """Create multiple relations between entities.""" logger.debug(f"Creating {len(relations_data)} relations") @@ -153,7 +153,7 @@ class MemoryService: logger.debug(f"Successfully created {len(relations)} relations") return relations - async def add_observations(self, observations_in: ObservationsIn) -> List[Observation]: + async def add_observations(self, observations_in: AddObservationsRequest) -> List[Observation]: """Add observations to an existing entity.""" logger.debug(f"Adding observations to entity: {observations_in.entity_id}") try: diff --git a/src/basic_memory/services/relation_service.py b/src/basic_memory/services/relation_service.py index 838ce05a..92861c05 100644 --- a/src/basic_memory/services/relation_service.py +++ b/src/basic_memory/services/relation_service.py @@ -1,9 +1,8 @@ """Service for managing relations in the database.""" from pathlib import Path -from basic_memory.models import Relation from basic_memory.repository.relation_repository import RelationRepository -from basic_memory.schemas import EntityIn, Relation +from basic_memory.schemas import EntityIn, RelationIn from . import DatabaseSyncError @@ -17,7 +16,7 @@ class RelationService: self.project_path = project_path self.relation_repo = relation_repo - async def create_relation(self, relation: Relation) -> Relation: + async def create_relation(self, relation: RelationIn) -> RelationIn: """Create a new relation in the database.""" try: return await self.relation_repo.create(relation.model_dump()) diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 9e9cfcaa..db39c958 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -186,8 +186,8 @@ async def test_add_observations(test_entity_data, memory_service, test_config): response = AddObservationsResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue] assert response.entity_id == entity_id - assert len(response.added_observations) == 1 - assert response.added_observations[0].content == "A new observation" + assert len(response.observations) == 1 + assert response.observations[0].content == "A new observation" @pytest.mark.anyio async def test_invalid_tool_name(test_config): diff --git a/tests/test_memory_service.py b/tests/test_memory_service.py index a41afac0..042a3e62 100644 --- a/tests/test_memory_service.py +++ b/tests/test_memory_service.py @@ -2,7 +2,7 @@ import pytest from basic_memory.services import MemoryService from basic_memory.fileio import read_entity_file -from basic_memory.schemas import CreateEntitiesInput, CreateRelationsInput, ObservationsIn, Relation +from basic_memory.schemas import CreateEntitiesRequest, CreateRelationsRequest, AddObservationsRequest, RelationIn test_entities_data = [ { @@ -21,7 +21,7 @@ test_entities_data = [ async def test_create_entities(memory_service: MemoryService): """Should create multiple entities in parallel with their observations.""" - entity_input = CreateEntitiesInput.model_validate({"entities": test_entities_data}) + entity_input = CreateEntitiesRequest.model_validate({"entities": test_entities_data}) entities = await memory_service.create_entities(entity_input.entities) # Verify the SQLAlchemy models were created @@ -51,7 +51,7 @@ 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.""" - entity_input = CreateEntitiesInput.model_validate({"entities": test_entities_data}) + entity_input = CreateEntitiesRequest.model_validate({"entities": test_entities_data}) entities = await memory_service.create_entities([entity_input.entities[0]]) entity = entities[0] @@ -65,7 +65,7 @@ async def test_add_observations(memory_service: MemoryService): } # Add observations - returns List[models.Observation] - observation_input = ObservationsIn.model_validate(observations_data) + observation_input = AddObservationsRequest.model_validate(observations_data) added_observations = await memory_service.add_observations(observation_input) # Check the SQLAlchemy model results @@ -94,13 +94,13 @@ 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 - observation_input = ObservationsIn.model_validate(observations_data) + observation_input = AddObservationsRequest.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.""" - entity_input = CreateEntitiesInput.model_validate({"entities": test_entities_data}) + entity_input = CreateEntitiesRequest.model_validate({"entities": test_entities_data}) entities = await memory_service.create_entities(entity_input.entities) entity1, entity2 = entities @@ -120,7 +120,7 @@ async def test_create_relations(memory_service: MemoryService): ] # Create relations - returns List[models.Relation] - input_args = CreateRelationsInput.model_validate({"relations": test_relations_data}) + input_args = CreateRelationsRequest.model_validate({"relations": test_relations_data}) relations = await memory_service.create_relations(input_args.relations) # Verify SQLAlchemy Relation models were created @@ -183,7 +183,7 @@ 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 - entity_input = CreateEntitiesInput.model_validate({"entities": test_entities_data}) + entity_input = CreateEntitiesRequest.model_validate({"entities": test_entities_data}) entities = await memory_service.create_entities([entity_input.entities[0]]) entity1 = entities[0] @@ -195,4 +195,4 @@ async def test_create_relations_with_invalid_entity_id(memory_service: MemorySer } with pytest.raises(Exception) as exc: # We might want to define a specific error type - await memory_service.create_relations([Relation.model_validate(bad_relation)]) \ No newline at end of file + await memory_service.create_relations([RelationIn.model_validate(bad_relation)]) \ No newline at end of file diff --git a/tests/test_relation_service.py b/tests/test_relation_service.py index 5f4221b8..96397edd 100644 --- a/tests/test_relation_service.py +++ b/tests/test_relation_service.py @@ -2,7 +2,7 @@ import pytest import pytest_asyncio -from basic_memory.schemas import EntityIn, Relation +from basic_memory.schemas import EntityIn, RelationIn pytestmark = pytest.mark.asyncio @@ -36,7 +36,7 @@ async def test_create_relation(relation_service, sample_entities): """Test creating a basic relation between two entities""" entity1, entity2 = sample_entities - relation_data = Relation( + relation_data = RelationIn( from_id=entity1.id, to_id=entity2.id, relation_type="test_relation" @@ -62,7 +62,7 @@ async def test_create_relation_with_context(relation_service, sample_entities): """Test creating a relation with context information""" entity1, entity2 = sample_entities - relation_data = Relation( + relation_data = RelationIn( from_id=entity1.id, to_id=entity2.id, relation_type="test_relation", diff --git a/tests/test_schemas.py b/tests/test_schemas.py index 4f930254..fe6ba017 100644 --- a/tests/test_schemas.py +++ b/tests/test_schemas.py @@ -4,10 +4,10 @@ from pydantic import ValidationError from basic_memory.schemas import ( EntityIn, EntityOut, - Relation, - CreateEntitiesInput, - SearchNodesInput, - OpenNodesInput, + RelationIn, + CreateEntitiesRequest, + SearchNodesRequest, + OpenNodesRequest, ) def test_entity_in_minimal(): @@ -67,7 +67,7 @@ def test_relation_in_validation(): "to_id": "456", "relation_type": "test" } - relation = Relation.model_validate(data) + relation = RelationIn.model_validate(data) assert relation.from_id == "123" assert relation.to_id == "456" assert relation.relation_type == "test" @@ -75,12 +75,12 @@ def test_relation_in_validation(): # With context data["context"] = "test context" - relation = Relation.model_validate(data) + relation = RelationIn.model_validate(data) assert relation.context == "test context" # Missing required fields with pytest.raises(ValidationError): - Relation.model_validate({"from_id": "123", "to_id": "456"}) # Missing relationType + RelationIn.model_validate({"from_id": "123", "to_id": "456"}) # Missing relationType def test_create_entities_input(): """Test CreateEntitiesInput validation.""" @@ -97,13 +97,13 @@ def test_create_entities_input(): } ] } - create_input = CreateEntitiesInput.model_validate(data) + create_input = CreateEntitiesRequest.model_validate(data) assert len(create_input.entities) == 2 assert create_input.entities[1].description == "test description" # Empty entities list should fail with pytest.raises(ValidationError): - CreateEntitiesInput.model_validate({"entities": []}) + CreateEntitiesRequest.model_validate({"entities": []}) def test_entity_out_from_attributes(): """Test EntityOut creation from database model attributes.""" @@ -167,17 +167,17 @@ def test_optional_fields(): def test_search_nodes_input(): """Test SearchNodesInput validation.""" - search = SearchNodesInput.model_validate({"query": "test query"}) + search = SearchNodesRequest.model_validate({"query": "test query"}) assert search.query == "test query" with pytest.raises(ValidationError): - SearchNodesInput.model_validate({}) # Missing required query + SearchNodesRequest.model_validate({}) # Missing required query def test_open_nodes_input(): """Test OpenNodesInput validation.""" - open_input = OpenNodesInput.model_validate({"names": ["entity1", "entity2"]}) + open_input = OpenNodesRequest.model_validate({"names": ["entity1", "entity2"]}) assert len(open_input.names) == 2 # Empty names list should fail with pytest.raises(ValidationError): - OpenNodesInput.model_validate({"names": []}) \ No newline at end of file + OpenNodesRequest.model_validate({"names": []}) \ No newline at end of file