From 6b684dcf9c5a9eaa1b7f2bed4976c72b0276e7f2 Mon Sep 17 00:00:00 2001 From: phernandez Date: Mon, 16 Dec 2024 20:43:01 -0600 Subject: [PATCH] sanitize inputs, fix tests --- src/basic_memory/api/routers/knowledge.py | 7 +- src/basic_memory/mcp/server.py | 2 +- src/basic_memory/schemas.py | 14 ++- src/basic_memory/services/memory_service.py | 129 ++++++++++++-------- tests/api/test_knowledge_router.py | 8 +- tests/mcp/test_create_relations.py | 3 +- tests/mcp/test_delete_entities.py | 5 +- tests/mcp/test_delete_relations.py | 11 +- tests/mcp/test_mcp_server.py | 2 +- tests/mcp/test_open_nodes.py | 7 +- tests/schemas/test_schemas.py | 90 ++++++-------- 11 files changed, 137 insertions(+), 141 deletions(-) diff --git a/src/basic_memory/api/routers/knowledge.py b/src/basic_memory/api/routers/knowledge.py index d5f4e9c7..e9f04d5d 100644 --- a/src/basic_memory/api/routers/knowledge.py +++ b/src/basic_memory/api/routers/knowledge.py @@ -24,6 +24,7 @@ from basic_memory.schemas import ( AddObservationsResponse, RelationResponse, DeleteEntityRequest, + Entity, ) router = APIRouter(prefix="/knowledge", tags=["knowledge"]) @@ -94,10 +95,8 @@ async def search_nodes( @router.post("/nodes", response_model=OpenNodesResponse) async def open_nodes(data: OpenNodesRequest, memory_service: MemoryServiceDep) -> OpenNodesResponse: """Open specific nodes by their names.""" - entities = await memory_service.open_nodes(data.names) - return OpenNodesResponse( - entities=[EntityResponse.model_validate(entity) for entity in entities] - ) + entities = await memory_service.open_nodes(data.entity_ids) + return OpenNodesResponse(entities=[Entity.model_validate(entity) for entity in entities]) ## Delete endpoints diff --git a/src/basic_memory/mcp/server.py b/src/basic_memory/mcp/server.py index fdd5fb2d..d8d56c87 100644 --- a/src/basic_memory/mcp/server.py +++ b/src/basic_memory/mcp/server.py @@ -173,7 +173,7 @@ class MemoryServer(Server): raise McpError(METHOD_NOT_FOUND, f"Unknown tool: {name}") # Make API call - logger.debug(f"Calling API endpoint for {name}") + logger.debug(f"Calling API endpoint for {name} with arguments: {arguments}") response = await handler(self.client, arguments) # Handle HTTP errors diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index c906f234..36b12332 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -5,14 +5,16 @@ from typing import List, Optional, Annotated from annotated_types import MinLen, MaxLen from pydantic import BaseModel, ConfigDict, BeforeValidator -from basic_memory.utils import sanitize_name - # Strip whitespace def strip_whitespace(obs: str) -> str: return obs.strip() +def lower_strip_whitespace(val: str) -> str: + return strip_whitespace(val.lower()) + + Observation = Annotated[str, BeforeValidator(strip_whitespace), MinLen(1), MaxLen(1000)] EntityType = Annotated[str, BeforeValidator(strip_whitespace), MinLen(1), MaxLen(20)] @@ -20,7 +22,7 @@ EntityType = Annotated[str, BeforeValidator(strip_whitespace), MinLen(1), MaxLen RelationType = Annotated[str, BeforeValidator(strip_whitespace), MinLen(1), MaxLen(20)] # Custom field types with validation -EntityId = Annotated[str, BeforeValidator(sanitize_name)] +EntityId = Annotated[str, BeforeValidator(lower_strip_whitespace)] class Relation(BaseModel): @@ -81,7 +83,7 @@ class SearchNodesRequest(BaseModel): class OpenNodesRequest(BaseModel): """Request schema for open_nodes tool.""" - names: Annotated[List[EntityId], MinLen(1)] + entity_ids: Annotated[List[str], MinLen(1)] class CreateRelationsRequest(BaseModel): @@ -165,9 +167,9 @@ class SearchNodesResponse(SQLAlchemyModel): class OpenNodesResponse(SQLAlchemyModel): - """Response for open_nodes tool.""" + """Response for open_nodes tool. This returns the Entity object because it is read from a file""" - entities: List[EntityResponse] + entities: List[Entity] class AddObservationsResponse(SQLAlchemyModel): diff --git a/src/basic_memory/services/memory_service.py b/src/basic_memory/services/memory_service.py index 9246d04c..2364665f 100644 --- a/src/basic_memory/services/memory_service.py +++ b/src/basic_memory/services/memory_service.py @@ -1,29 +1,40 @@ """Service for orchestrating entity, relation, and observation operations.""" -import asyncio -from typing import List, Dict, Any, Optional, Sequence -from pathlib import Path -from basic_memory.models import Entity as EntityModel, Observation as ObservationModel, Relation as RelationModel -from basic_memory.schemas import ( - AddObservationsRequest, Entity, Relation -) -from basic_memory.fileio import write_entity_file, read_entity_file, EntityNotFoundError, get_entity_path -from basic_memory.services import EntityService, RelationService, ObservationService +import asyncio +from pathlib import Path +from typing import List, Dict, Any, Optional, Sequence + from loguru import logger +from basic_memory.fileio import ( + write_entity_file, + read_entity_file, + EntityNotFoundError, + get_entity_path, +) +from basic_memory.models import ( + Entity as EntityModel, + Observation as ObservationModel, + Relation as RelationModel, +) +from basic_memory.schemas import AddObservationsRequest, Entity, Relation +from basic_memory.services import EntityService, RelationService, ObservationService + class MemoryService: """Orchestrates entity, relation, and observation operations with filesystem handling.""" - + def __init__( self, project_path: Optional[Path], entity_service: EntityService, relation_service: RelationService, - observation_service: ObservationService + observation_service: ObservationService, ): if project_path: - assert project_path.is_dir(), "Path does not exist or is not a directory: {project_path}" + assert ( + project_path.is_dir() + ), "Path does not exist or is not a directory: {project_path}" self.project_path = project_path self.entities_path = project_path / "entities" @@ -40,15 +51,15 @@ class MemoryService: logger.debug(f"Creating {len(entities_in)} entities") # TODO this could be better - for entity in entities_in: + for e in entities_in: + # Generate ID and write file try: - existing = await self.entity_service.get_by_type_and_name( - entity.entity_type, - entity.name + existing = await self.entity_service.get_entity( + EntityModel.generate_id(e.entity_type, e.name) ) if existing: raise ValueError( - f"Entity already exists: {entity.entity_type}/{entity.name}" + f"Entity {e.entity_type}/{e.name} already exists, id: {EntityModel.generate_id(e.entity_type, e.name)}" ) except EntityNotFoundError: # Good - entity doesn't exist yet @@ -65,28 +76,34 @@ class MemoryService: await asyncio.gather(*file_writes) logger.debug("Completed all file writes") - async def create_entity_in_db(entity_in: Entity): - logger.debug(f"Creating entity in DB: {entity_in}") + async def create_entity_in_db(entity_create: Entity): + logger.debug(f"Creating entity in DB: {entity_create}") try: # Create base entity - created_entity = await self.entity_service.create_entity(entity_in) + created_entity = await self.entity_service.create_entity(entity_create) logger.debug(f"Created base entity: {created_entity.id}") # Add observations - await self.observation_service.add_observations(created_entity.id, entity_in.observations) - logger.debug(f"Added {len(entity_in.observations)} observations to {created_entity.id}") + await self.observation_service.add_observations( + created_entity.id, entity_create.observations + ) + logger.debug( + f"Added {len(entity_create.observations)} observations to {created_entity.id}" + ) # Add relations - for relation in entity_in.relations: + for relation in entity_create.relations: await self.relation_service.create_relation(relation) - logger.debug(f"Added {len(entity_in.relations)} relations for {created_entity.id}") + logger.debug( + f"Added {len(entity_create.relations)} relations for {created_entity.id}" + ) # Query final state final_entity = await self.entity_service.get_entity(created_entity.id) logger.debug(f"Retrieved final entity state: {final_entity}") return final_entity except Exception: - logger.exception(f"Failed to create entity in DB: {entity_in}") + logger.exception(f"Failed to create entity in DB: {entity_create}") raise # Update database index sequentially @@ -126,12 +143,12 @@ class MemoryService: logger.debug(f"Processing relation: {relation.from_id} -> {relation.to_id}") try: # First read complete entities from filesystem - from_entity = await read_entity_file(self.entities_path, relation.from_id) + from_entity = await read_entity_file(self.entities_path, relation.from_id) to_entity = await read_entity_file(self.entities_path, relation.to_id) logger.debug(f"Read entities for relation: {from_entity.id}, {to_entity.id}") # Add the new relation to the source entity - if not hasattr(from_entity, 'relations'): + if not hasattr(from_entity, "relations"): from_entity.relations = [] from_entity.relations.append(relation) logger.debug(f"Added relation to source entity: {from_entity.id}") @@ -142,8 +159,10 @@ class MemoryService: assert to_entity.id is not None await asyncio.gather( - *[write_entity_file(self.entities_path, from_entity.id, from_entity), - write_entity_file(self.entities_path, to_entity.id, to_entity)] + *[ + write_entity_file(self.entities_path, from_entity.id, from_entity), + write_entity_file(self.entities_path, to_entity.id, to_entity), + ] ) logger.debug("Wrote updated entity files") @@ -158,14 +177,16 @@ class MemoryService: logger.debug(f"Successfully created {len(relations)} relations") return relations - async def add_observations(self, observations_in: AddObservationsRequest) -> List[ObservationModel]: + async def add_observations( + self, observations_in: AddObservationsRequest + ) -> List[ObservationModel]: """Add observations to an existing entity.""" logger.debug(f"Adding observations to entity: {observations_in.entity_id}") try: # First get the entity from DB to get its ID db_entity = await self.entity_service.get_entity(observations_in.entity_id) logger.debug(f"Found entity in DB: {db_entity.id}") - + # Read entity from filesystem using the ID entity = await read_entity_file(self.entities_path, db_entity.id) logger.debug(f"Read entity from filesystem: {db_entity.id}") @@ -178,9 +199,11 @@ class MemoryService: logger.debug("Writing updated entity file") await write_entity_file(self.entities_path, db_entity.id, entity) logger.debug("Wrote updated entity file") - + # Update database index - added_observations = await self.observation_service.add_observations(db_entity.id, observations_in.observations) + added_observations = await self.observation_service.add_observations( + db_entity.id, observations_in.observations + ) logger.debug(f"Added {len(added_observations)} observations to DB") return added_observations @@ -206,7 +229,7 @@ class MemoryService: if file_path.exists(): file_path.unlink() logger.debug(f"Deleted entity {entity_id} file: {file_path}") - + # Then update database result = await self.entity_service.delete_entity(entity_id) if result: @@ -224,20 +247,17 @@ class MemoryService: try: # First read the entity entity = await read_entity_file(self.entities_path, entity_id) - + # Remove observations from entity original_count = len(entity.observations) - entity.observations = [ - obs for obs in entity.observations - if obs not in contents - ] - + entity.observations = [obs for obs in entity.observations if obs not in contents] + # Only write file if we actually removed anything if len(entity.observations) < original_count: # Write updated entity file first (source of truth) await write_entity_file(self.entities_path, entity_id, entity) logger.debug(f"Updated entity file: {entity_id}") - + # Then update database result = await self.observation_service.delete_observations(entity_id, contents) logger.debug(f"Deleted observations from database: {result}") @@ -256,25 +276,27 @@ class MemoryService: for relation in relations: # First read the source entity try: - from_entity = await read_entity_file(self.entities_path, relation['from_id']) + from_entity = await read_entity_file(self.entities_path, relation["from_id"]) except EntityNotFoundError: logger.debug(f"Source entity not found: {relation['from_id']}") continue # Update the entity's relations - if hasattr(from_entity, 'relations'): + if hasattr(from_entity, "relations"): original_count = len(from_entity.relations) - relation_type = relation.get('relation_type') - + relation_type = relation.get("relation_type") + if relation_type: from_entity.relations = [ - r for r in from_entity.relations - if not (r.to_id == relation['to_id'] and r.relation_type == relation_type) + r + for r in from_entity.relations + if not ( + r.to_id == relation["to_id"] and r.relation_type == relation_type + ) ] else: from_entity.relations = [ - r for r in from_entity.relations - if r.to_id != relation['to_id'] + r for r in from_entity.relations if r.to_id != relation["to_id"] ] # Only write file if we actually removed any relations @@ -337,10 +359,13 @@ class MemoryService: return None try: - entities = [entity for entity in await asyncio.gather(*(read_node(id) for id in entity_ids)) - if entity is not None] + entities = [ + entity + for entity in await asyncio.gather(*(read_node(id) for id in entity_ids)) + if entity is not None + ] logger.debug(f"Opened {len(entities)} entities") return entities except Exception: logger.exception("Failed to open nodes") - raise \ No newline at end of file + raise diff --git a/tests/api/test_knowledge_router.py b/tests/api/test_knowledge_router.py index 1b9aa949..2835dbbf 100644 --- a/tests/api/test_knowledge_router.py +++ b/tests/api/test_knowledge_router.py @@ -257,9 +257,7 @@ async def test_delete_observations(client, observation_repository): # Delete specific observations request_data = {"entity_id": entity.id, "deletions": [observations[0].content]} - response = await client.post( - f"/knowledge/entities/{entity.id}/observations/delete", json=request_data - ) + response = await client.post("/knowledge/observations/delete", json=request_data) assert response.status_code == 200 assert response.json() == {"deleted": True} @@ -312,9 +310,7 @@ async def test_delete_nonexistent_observations(client): entity = await create_entity(client) request_data = {"entity_id": entity.id, "deletions": ["Nonexistent observation"]} - response = await client.post( - f"/knowledge/entities/{entity.id}/observations/delete", json=request_data - ) + response = await client.post("/knowledge/observations/delete", json=request_data) assert response.status_code == 200 assert response.json() == {"deleted": False} diff --git a/tests/mcp/test_create_relations.py b/tests/mcp/test_create_relations.py index 3278e2b7..2cab222e 100644 --- a/tests/mcp/test_create_relations.py +++ b/tests/mcp/test_create_relations.py @@ -3,7 +3,6 @@ import pytest from basic_memory.schemas import SearchNodesResponse -from basic_memory.utils import sanitize_name @pytest.mark.asyncio @@ -39,5 +38,5 @@ async def test_create_relations(test_entity_data, client, server): assert len(response.matches) == 1 entity = response.matches[0] assert len(entity.relations) == 1 - assert entity.relations[0].to_id == sanitize_name("test/TestEntityB") + assert entity.relations[0].to_id == "test/testentityb" assert entity.relations[0].relation_type == "relates_to" diff --git a/tests/mcp/test_delete_entities.py b/tests/mcp/test_delete_entities.py index 08f6442f..d1d72cdb 100644 --- a/tests/mcp/test_delete_entities.py +++ b/tests/mcp/test_delete_entities.py @@ -3,7 +3,6 @@ import pytest from basic_memory.schemas import SearchNodesResponse -from basic_memory.utils import sanitize_name @pytest.mark.asyncio @@ -19,9 +18,7 @@ async def test_delete_entities(server): await server.handle_call_tool("create_entities", entities) # Delete first entity - await server.handle_call_tool( - "delete_entities", {"entity_ids": [sanitize_name("test/DeleteTest1")]} - ) + await server.handle_call_tool("delete_entities", {"entity_ids": ["test/deletetest1"]}) # Verify through search search_result = await server.handle_call_tool("search_nodes", {"query": "DeleteTest"}) diff --git a/tests/mcp/test_delete_relations.py b/tests/mcp/test_delete_relations.py index bbaf8c20..694dcee8 100644 --- a/tests/mcp/test_delete_relations.py +++ b/tests/mcp/test_delete_relations.py @@ -3,7 +3,6 @@ import pytest from basic_memory.schemas import SearchNodesResponse -from basic_memory.utils import sanitize_name @pytest.mark.asyncio @@ -22,8 +21,8 @@ async def test_delete_relations(server): relation = { "relations": [ { - "from_id": sanitize_name("test/RelSource"), - "to_id": sanitize_name("test/RelTarget"), + "from_id": "test/relsource", + "to_id": "test/reltarget", "relation_type": "relates_to", } ] @@ -36,8 +35,8 @@ async def test_delete_relations(server): { "relations": [ { - "from_id": sanitize_name("test/RelSource"), - "to_id": sanitize_name("test/RelTarget"), + "from_id": "test/relsource", + "to_id": "test/reltarget", "relation_type": "relates_to", } ] @@ -45,7 +44,7 @@ async def test_delete_relations(server): ) # Verify through search - search_result = await server.handle_call_tool("search_nodes", {"query": "RelSource"}) + search_result = await server.handle_call_tool("search_nodes", {"query": "relsource"}) search_response = SearchNodesResponse.model_validate_json(search_result[0].resource.text) # Source entity should exist but have no relations diff --git a/tests/mcp/test_mcp_server.py b/tests/mcp/test_mcp_server.py index 935daa83..bff52e17 100644 --- a/tests/mcp/test_mcp_server.py +++ b/tests/mcp/test_mcp_server.py @@ -35,7 +35,7 @@ async def test_empty_arrays(server): assert INVALID_PARAMS == exc.value.args[0] with pytest.raises(McpError) as exc: - await server.handle_call_tool("open_nodes", {"names": []}) + await server.handle_call_tool("open_nodes", {"entity_ids": []}) assert INVALID_PARAMS == exc.value.args[0] diff --git a/tests/mcp/test_open_nodes.py b/tests/mcp/test_open_nodes.py index 30ea4ec7..4f77d444 100644 --- a/tests/mcp/test_open_nodes.py +++ b/tests/mcp/test_open_nodes.py @@ -5,7 +5,6 @@ from mcp.types import EmbeddedResource from basic_memory.mcp.server import MIME_TYPE from basic_memory.schemas import OpenNodesResponse -from basic_memory.utils import sanitize_name @pytest.mark.asyncio @@ -36,7 +35,7 @@ async def test_open_nodes(server): # Open specific nodes result = await server.handle_call_tool( - "open_nodes", {"names": ["test/opentesta", "test/opentestb"]} + "open_nodes", {"entity_ids": ["test/opentesta", "test/opentestb"]} ) # Verify response format @@ -55,7 +54,7 @@ async def test_open_nodes(server): # Verify entity content entity = response.entities[0] - assert entity.id == sanitize_name("test/OpenTestA") + assert entity.id == "test/opentesta" assert entity.entity_type == "test" assert len(entity.observations) == 1 - assert entity.observations[0].content == "First test entity" + assert entity.observations[0] == "First test entity" diff --git a/tests/schemas/test_schemas.py b/tests/schemas/test_schemas.py index f52b51e8..8962b9b6 100644 --- a/tests/schemas/test_schemas.py +++ b/tests/schemas/test_schemas.py @@ -1,6 +1,8 @@ """Tests for Pydantic schema validation and conversion.""" + import pytest from pydantic import ValidationError + from basic_memory.schemas import ( Entity, EntityResponse, @@ -10,12 +12,10 @@ from basic_memory.schemas import ( OpenNodesRequest, ) + def test_entity_in_minimal(): """Test creating EntityIn with minimal required fields.""" - data = { - "name": "test_entity", - "entity_type": "test" - } + data = {"name": "test_entity", "entity_type": "test"} entity = Entity.model_validate(data) assert entity.name == "test_entity" assert entity.entity_type == "test" @@ -23,22 +23,15 @@ def test_entity_in_minimal(): assert entity.observations == [] assert entity.relations == [] + def test_entity_in_complete(): """Test creating EntityIn with all fields.""" data = { "name": "test_entity", "entity_type": "test", "description": "A test entity", - "observations": [ - "Test observation" - ], - "relations": [ - { - "from_id": "123", - "to_id": "456", - "relation_type": "test_relation" - } - ] + "observations": ["Test observation"], + "relations": [{"from_id": "123", "to_id": "456", "relation_type": "test_relation"}], } entity = Entity.model_validate(data) assert entity.name == "test_entity" @@ -49,6 +42,7 @@ def test_entity_in_complete(): assert len(entity.relations) == 1 assert entity.relations[0].from_id == "123" + def test_entity_in_validation(): """Test validation errors for EntityIn.""" with pytest.raises(ValidationError): @@ -60,13 +54,10 @@ def test_entity_in_validation(): with pytest.raises(ValidationError): Entity.model_validate({"entityType": "test"}) # Missing name + def test_relation_in_validation(): """Test RelationIn validation.""" - data = { - "from_id": "123", - "to_id": "456", - "relation_type": "test" - } + data = {"from_id": "123", "to_id": "456", "relation_type": "test"} relation = Relation.model_validate(data) assert relation.from_id == "123" assert relation.to_id == "456" @@ -82,19 +73,13 @@ def test_relation_in_validation(): with pytest.raises(ValidationError): Relation.model_validate({"from_id": "123", "to_id": "456"}) # Missing relationType + def test_create_entities_input(): """Test CreateEntitiesInput validation.""" data = { "entities": [ - { - "name": "entity1", - "entity_type": "test" - }, - { - "name": "entity2", - "entity_type": "test", - "description": "test description" - } + {"name": "entity1", "entity_type": "test"}, + {"name": "entity2", "entity_type": "test", "description": "test description"}, ] } create_input = CreateEntityRequest.model_validate(data) @@ -105,6 +90,7 @@ def test_create_entities_input(): with pytest.raises(ValidationError): CreateEntityRequest.model_validate({"entities": []}) + def test_entity_out_from_attributes(): """Test EntityOut creation from database model attributes.""" # Simulate database model attributes @@ -113,18 +99,10 @@ def test_entity_out_from_attributes(): "name": "test", "entity_type": "test", "description": "test description", - "observations": [ - {"id": 1, "content": "test obs", "context": None} - ], + "observations": [{"id": 1, "content": "test obs", "context": None}], "relations": [ - { - "id": 1, - "from_id": "123", - "to_id": "456", - "relation_type": "test", - "context": None - } - ] + {"id": 1, "from_id": "123", "to_id": "456", "relation_type": "test", "context": None} + ], } entity = EntityResponse.model_validate(db_data) assert entity.id == "123" @@ -133,6 +111,7 @@ def test_entity_out_from_attributes(): assert entity.observations[0].id == 1 assert len(entity.relations) == 1 + def test_optional_fields(): """Test handling of optional fields.""" # Create with no optional fields @@ -142,28 +121,28 @@ def test_optional_fields(): assert entity.relations == [] # Create with empty optional fields - entity = Entity.model_validate({ - "name": "test", - "entity_type": "test", - "description": None, - "observations": [], - "relations": [] - }) + entity = Entity.model_validate( + { + "name": "test", + "entity_type": "test", + "description": None, + "observations": [], + "relations": [], + } + ) assert entity.description is None assert entity.observations == [] assert entity.relations == [] # Create with some optional fields - entity = Entity.model_validate({ - "name": "test", - "entity_type": "test", - "description": "test", - "observations": [] - }) + entity = Entity.model_validate( + {"name": "test", "entity_type": "test", "description": "test", "observations": []} + ) assert entity.description == "test" assert entity.observations == [] assert entity.relations == [] + def test_search_nodes_input(): """Test SearchNodesInput validation.""" search = SearchNodesRequest.model_validate({"query": "test query"}) @@ -172,11 +151,12 @@ def test_search_nodes_input(): with pytest.raises(ValidationError): SearchNodesRequest.model_validate({}) # Missing required query + def test_open_nodes_input(): """Test OpenNodesInput validation.""" - open_input = OpenNodesRequest.model_validate({"names": ["entity1", "entity2"]}) - assert len(open_input.names) == 2 + open_input = OpenNodesRequest.model_validate({"entity_ids": ["entity1", "entity2"]}) + assert len(open_input.entity_ids) == 2 # Empty names list should fail with pytest.raises(ValidationError): - OpenNodesRequest.model_validate({"names": []}) \ No newline at end of file + OpenNodesRequest.model_validate({"entity_ids": []})