diff --git a/src/basic_memory/fileio.py b/src/basic_memory/fileio.py index 612c71a2..9371f7f1 100644 --- a/src/basic_memory/fileio.py +++ b/src/basic_memory/fileio.py @@ -176,17 +176,17 @@ 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(RelationIn( - from_id=entity_id, - to_id=target_id, - relation_type=relation_type, + 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] context=context )) - return EntityIn( + return EntityIn( # pyright: ignore [reportCallIssue] id=entity_id, name=name, - entity_type=entity_type, + entity_type=entity_type, # pyright: ignore [reportCallIssue] observations=observations, relations=relations ) diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index 84a8bcc3..9459c3f8 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -97,6 +97,12 @@ 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] @@ -123,6 +129,10 @@ class OpenNodesResponse(SQLAlchemyOut): """Response for open_nodes tool.""" entities: List[EntityOut] +class AddObservationsResponse(SQLAlchemyOut): + """Response for add_observations tool.""" + entity_id: str + added_observations: List[ObservationOut] class CreateRelationsResponse(SQLAlchemyOut): """Response for create_relations tool.""" diff --git a/src/basic_memory/services/entity_service.py b/src/basic_memory/services/entity_service.py index bde20722..9fa3f1b5 100644 --- a/src/basic_memory/services/entity_service.py +++ b/src/basic_memory/services/entity_service.py @@ -96,6 +96,9 @@ class EntityService: logger.exception(f"Failed to get entity by type/name: {entity_type}/{name}") raise + async def get_all(self) -> Sequence[Entity]: + return await self.entity_repo.find_all() + async def delete_entity(self, entity_id: str) -> bool: """Delete entity from database.""" logger.debug(f"Deleting entity: {entity_id}") @@ -105,4 +108,5 @@ class EntityService: return result except Exception: logger.exception(f"Failed to delete entity: {entity_id}") - raise \ No newline at end of file + raise + diff --git a/src/basic_memory/services/memory_service.py b/src/basic_memory/services/memory_service.py index f98aeb61..091b0476 100644 --- a/src/basic_memory/services/memory_service.py +++ b/src/basic_memory/services/memory_service.py @@ -133,9 +133,12 @@ class MemoryService: # Write updated entity files (filesystem is source of truth) logger.debug("Writing updated entity files") + assert from_entity.id is not None + 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") diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 927ad51f..e26c675f 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -1,14 +1,11 @@ """Tests for the MCP server implementation.""" import pytest -from pathlib import Path 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 ( - CreateEntitiesResponse, SearchNodesResponse, OpenNodesResponse, - AddObservationsResponse + CreateEntitiesResponse, SearchNodesResponse, AddObservationsResponse, ) @pytest.fixture diff --git a/tests/test_memory_service.py b/tests/test_memory_service.py index 27ddb810..dd73a692 100644 --- a/tests/test_memory_service.py +++ b/tests/test_memory_service.py @@ -3,7 +3,7 @@ import pytest from basic_memory.services import MemoryService from basic_memory.fileio import read_entity_file from basic_memory.models import Entity as EntityModel, Observation, Relation -from basic_memory.schemas import EntityIn, CreateEntitiesInput, AddObservationsInput, CreateRelationsInput +from basic_memory.schemas import EntityIn, CreateEntitiesInput, CreateRelationsInput, ObservationsIn, RelationIn test_entities_data = [ { @@ -66,7 +66,7 @@ async def test_add_observations(memory_service: MemoryService): } # Add observations - returns List[models.Observation] - observation_input = AddObservationsInput.model_validate(observations_data) + observation_input = ObservationsIn.model_validate(observations_data) added_observations = await memory_service.add_observations(observation_input) # Check the SQLAlchemy model results @@ -96,7 +96,7 @@ 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 = AddObservationsInput.model_validate(observations_data) + observation_input = ObservationsIn.model_validate(observations_data) await memory_service.add_observations(observation_input) @pytest.mark.asyncio @@ -197,4 +197,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([bad_relation]) \ No newline at end of file + await memory_service.create_relations([RelationIn.model_validate(bad_relation)]) \ No newline at end of file