From b471e1f8d7e1ae7ae4594cec0a20a8f427e4e45c Mon Sep 17 00:00:00 2001 From: phernandez Date: Sat, 14 Dec 2024 19:30:01 -0600 Subject: [PATCH] fix schema plural namings --- src/basic_memory/api/routers/knowledge.py | 22 +++++++++++----------- src/basic_memory/mcp/server.py | 20 ++++++++++---------- src/basic_memory/schemas.py | 22 +++++++++++----------- tests/test_mcp_server.py | 8 ++++---- tests/test_memory_service.py | 10 +++++----- tests/test_schemas.py | 6 +++--- 6 files changed, 44 insertions(+), 44 deletions(-) diff --git a/src/basic_memory/api/routers/knowledge.py b/src/basic_memory/api/routers/knowledge.py index bbbc9645..dc6d62da 100644 --- a/src/basic_memory/api/routers/knowledge.py +++ b/src/basic_memory/api/routers/knowledge.py @@ -4,26 +4,26 @@ from fastapi import APIRouter from basic_memory.deps import MemoryServiceDep from basic_memory.schemas import ( - CreateEntitiesRequest, CreateEntitiesResponse, + CreateEntityRequest, CreateEntityResponse, SearchNodesRequest, SearchNodesResponse, CreateRelationsRequest, CreateRelationsResponse, EntityResponse, RelationResponse, AddObservationsRequest, ObservationResponse, OpenNodesRequest, OpenNodesResponse, - DeleteEntitiesResponse, + DeleteEntityResponse, DeleteObservationsRequest, DeleteObservationsResponse, AddObservationsResponse ) router = APIRouter(prefix="/knowledge", tags=["knowledge"]) -@router.post("/entities", response_model=CreateEntitiesResponse) +@router.post("/entities", response_model=CreateEntityResponse) async def create_entities( - data: CreateEntitiesRequest, + data: CreateEntityRequest, memory_service: MemoryServiceDep -) -> CreateEntitiesResponse: +) -> CreateEntityResponse: """Create new entities in the knowledge graph.""" entities = await memory_service.create_entities(data.entities) - return CreateEntitiesResponse(entities=[EntityResponse.model_validate(entity) for entity in entities]) + return CreateEntityResponse(entities=[EntityResponse.model_validate(entity) for entity in entities]) @router.get("/entities/{entity_id:path}", response_model=EntityResponse) @@ -36,14 +36,14 @@ async def get_entity( return EntityResponse.model_validate(entity) -@router.delete("/entities/{entity_id}", response_model=DeleteEntitiesResponse) +@router.delete("/entities/{entity_id}", response_model=DeleteEntityResponse) async def delete_entity( entity_id: str, memory_service: MemoryServiceDep -) -> DeleteEntitiesResponse: +) -> DeleteEntityResponse: """Delete a specific entity by ID.""" deleted = await memory_service.delete_entities([entity_id]) - return DeleteEntitiesResponse(deleted=deleted) # pyright: ignore [reportArgumentType] + return DeleteEntityResponse(deleted=deleted) # pyright: ignore [reportArgumentType] @router.post("/nodes", response_model=OpenNodesResponse) @@ -66,11 +66,11 @@ async def create_relations( return CreateRelationsResponse(relations=[RelationResponse.model_validate(relation) for relation in relations]) -@router.delete("/relations/{relation_id}", response_model=DeleteEntitiesResponse) +@router.delete("/relations/{relation_id}", response_model=DeleteEntityResponse) async def delete_relation( relation_id: int, memory_service: MemoryServiceDep -) -> DeleteEntitiesResponse: +) -> DeleteEntityResponse: """Delete a specific relation by ID.""" # TODO: Implement delete_relation in memory service raise NotImplementedError("Delete relation not implemented yet") diff --git a/src/basic_memory/mcp/server.py b/src/basic_memory/mcp/server.py index 4214d44a..98a5495b 100644 --- a/src/basic_memory/mcp/server.py +++ b/src/basic_memory/mcp/server.py @@ -21,12 +21,12 @@ from basic_memory.repository.observation_repository import ObservationRepository from basic_memory.repository.relation_repository import RelationRepository from basic_memory.schemas import ( # Tool inputs - CreateEntitiesRequest, SearchNodesRequest, OpenNodesRequest, - CreateRelationsRequest, DeleteEntitiesRequest, + CreateEntityRequest, SearchNodesRequest, OpenNodesRequest, + CreateRelationsRequest, DeleteEntityRequest, DeleteObservationsRequest, # Tool responses - CreateEntitiesResponse, SearchNodesResponse, OpenNodesResponse, - AddObservationsResponse, CreateRelationsResponse, DeleteEntitiesResponse, + CreateEntityResponse, SearchNodesResponse, OpenNodesResponse, + AddObservationsResponse, CreateRelationsResponse, DeleteEntityResponse, EntityResponse, ObservationResponse, RelationResponse, AddObservationsRequest ) from basic_memory.services import EntityService, ObservationService, RelationService @@ -103,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 = CreateEntitiesRequest.model_validate(args) + input_args = CreateEntityRequest.model_validate(args) logger.debug(f"Validated input: {len(input_args.entities)} entities") # Call service with validated data @@ -111,7 +111,7 @@ async def handle_create_entities( logger.debug(f"Created {len(entities)} entities") # Format response - response = CreateEntitiesResponse(entities=[EntityResponse.model_validate(entity) for entity in entities]) + response = CreateEntityResponse(entities=[EntityResponse.model_validate(entity) for entity in entities]) logger.debug("Formatted create_entities response") return create_response(response) @@ -192,10 +192,10 @@ async def handle_delete_entities( ) -> EmbeddedResource: """Handle delete_entities tool call.""" logger.debug(f"Deleting entities: {args}") - input_args = DeleteEntitiesRequest.model_validate(args) + input_args = DeleteEntityRequest.model_validate(args) deleted = await service.delete_entities(input_args.names) logger.debug(f"Deleted entities: {deleted}") - response = DeleteEntitiesResponse(deleted=deleted) + response = DeleteEntityResponse(deleted=deleted) return create_response(response) @@ -240,7 +240,7 @@ class MemoryServer(Server): Tool( name="create_entities", description="Create multiple new entities in the knowledge graph", - inputSchema=CreateEntitiesRequest.model_json_schema() + inputSchema=CreateEntityRequest.model_json_schema() ), Tool( name="search_nodes", @@ -265,7 +265,7 @@ class MemoryServer(Server): Tool( name="delete_entities", description="Delete entities from the knowledge graph", - inputSchema=DeleteEntitiesRequest.model_json_schema() + inputSchema=DeleteEntityRequest.model_json_schema() ), Tool( name="delete_observations", diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index e2afc7b8..49d97a5d 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -76,33 +76,33 @@ class EntityResponse(EntityBase, SQLAlchemyModel): relations: List[RelationResponse] = [] model_config = ConfigDict(populate_by_name=True) -# Tool Input Schemas -class CreateEntitiesRequest(BaseModel): - """Input schema for create_entities tool.""" +# Tool Request schemas +class CreateEntityRequest(BaseModel): + """Request schema for create_entities tool.""" entities: Annotated[List[EntityRequest], Len(min_length=1)] class SearchNodesRequest(BaseModel): - """Input schema for search_nodes tool.""" + """Request schema for search_nodes tool.""" query: str class OpenNodesRequest(BaseModel): - """Input schema for open_nodes tool.""" + """Request schema for open_nodes tool.""" names: Annotated[List[str], Len(min_length=1)] class CreateRelationsRequest(BaseModel): - """Input schema for create_relations tool.""" + """Request schema for create_relations tool.""" relations: List[RelationRequest] -class DeleteEntitiesRequest(BaseModel): - """Input schema for delete_entities tool.""" +class DeleteEntityRequest(BaseModel): + """Request schema for delete_entities tool.""" names: List[str] class DeleteObservationsRequest(BaseModel): - """Input schema for delete_observations tool.""" + """Request schema for delete_observations tool.""" entity_id: str deletions: List[str] # TODO: Make this more specific -class CreateEntitiesResponse(SQLAlchemyModel): +class CreateEntityResponse(SQLAlchemyModel): """Response for create_entities tool.""" entities: List[EntityResponse] @@ -124,7 +124,7 @@ class CreateRelationsResponse(SQLAlchemyModel): """Response for create_relations tool.""" relations: List[RelationResponse] -class DeleteEntitiesResponse(SQLAlchemyModel): +class DeleteEntityResponse(SQLAlchemyModel): """Response for delete_entities tool.""" deleted: List[str] diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index db39c958..774a36a4 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -5,7 +5,7 @@ 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.schemas import ( - CreateEntitiesResponse, SearchNodesResponse, AddObservationsResponse, + CreateEntityResponse, SearchNodesResponse, AddObservationsResponse, ) @pytest.fixture @@ -90,7 +90,7 @@ async def test_create_directory_entity(test_directory_entity_data, memory_servic assert result[0].type == "resource" # Verify entity creation - response = CreateEntitiesResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue] + response = CreateEntityResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue] assert len(response.entities) == 1 assert response.entities[0].name == "Directory Organization" assert response.entities[0].entity_type == "memory" @@ -115,7 +115,7 @@ async def test_create_entities_snake_case(test_entity_snake_case, memory_service assert str(result[0].resource.uri) == str(BASIC_MEMORY_URI) assert result[0].resource.mimeType == MIME_TYPE - response = CreateEntitiesResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue] + response = CreateEntityResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue] assert len(response.entities) == 1 assert response.entities[0].name == "Test Entity" assert response.entities[0].entity_type == "test" @@ -164,7 +164,7 @@ async def test_add_observations(test_entity_data, memory_service, test_config): memory_service=memory_service ) - create_response = CreateEntitiesResponse.model_validate_json(create_result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue] + create_response = CreateEntityResponse.model_validate_json(create_result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue] entity_id = create_response.entities[0].id # Add new observations using camelCase diff --git a/tests/test_memory_service.py b/tests/test_memory_service.py index 83302f87..ed3baeb2 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 CreateEntitiesRequest, CreateRelationsRequest, AddObservationsRequest, RelationRequest +from basic_memory.schemas import CreateEntityRequest, CreateRelationsRequest, AddObservationsRequest, RelationRequest 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 = CreateEntitiesRequest.model_validate({"entities": test_entities_data}) + entity_input = CreateEntityRequest.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 = CreateEntitiesRequest.model_validate({"entities": test_entities_data}) + entity_input = CreateEntityRequest.model_validate({"entities": test_entities_data}) entities = await memory_service.create_entities([entity_input.entities[0]]) entity = entities[0] @@ -100,7 +100,7 @@ async def test_add_observations_nonexistent_entity(memory_service: MemoryService @pytest.mark.asyncio async def test_create_relations(memory_service: MemoryService): """Should create relations between entities and update both filesystem and database.""" - entity_input = CreateEntitiesRequest.model_validate({"entities": test_entities_data}) + entity_input = CreateEntityRequest.model_validate({"entities": test_entities_data}) entities = await memory_service.create_entities(entity_input.entities) entity1, entity2 = entities @@ -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 = CreateEntitiesRequest.model_validate({"entities": test_entities_data}) + entity_input = CreateEntityRequest.model_validate({"entities": test_entities_data}) entities = await memory_service.create_entities([entity_input.entities[0]]) entity1 = entities[0] diff --git a/tests/test_schemas.py b/tests/test_schemas.py index 87f4e682..fa5f72ea 100644 --- a/tests/test_schemas.py +++ b/tests/test_schemas.py @@ -5,7 +5,7 @@ from basic_memory.schemas import ( EntityRequest, EntityResponse, RelationRequest, - CreateEntitiesRequest, + CreateEntityRequest, SearchNodesRequest, OpenNodesRequest, ) @@ -97,13 +97,13 @@ def test_create_entities_input(): } ] } - create_input = CreateEntitiesRequest.model_validate(data) + create_input = CreateEntityRequest.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): - CreateEntitiesRequest.model_validate({"entities": []}) + CreateEntityRequest.model_validate({"entities": []}) def test_entity_out_from_attributes(): """Test EntityOut creation from database model attributes."""