From 86329e0105f68703f8ec9c03d47256da72fa4436 Mon Sep 17 00:00:00 2001 From: phernandez Date: Mon, 16 Dec 2024 19:35:18 -0600 Subject: [PATCH] sanitize inputs for id/name --- src/basic_memory/api/routers/knowledge.py | 7 +- src/basic_memory/mcp/server.py | 8 +- src/basic_memory/models.py | 75 ++-- src/basic_memory/schemas.py | 37 +- src/basic_memory/utils.py | 34 +- tests/mcp/conftest.py | 66 +++ tests/mcp/test_add_observations.py | 39 ++ tests/mcp/test_create_entitites.py | 107 +++++ tests/mcp/test_create_relations.py | 43 ++ tests/mcp/test_delete_entities.py | 32 ++ tests/mcp/test_delete_observations.py | 51 +++ tests/mcp/test_delete_relations.py | 53 +++ tests/mcp/test_list_tools.py | 30 ++ tests/mcp/test_mcp_server.py | 513 ++++++++-------------- tests/mcp/test_open_nodes.py | 61 +++ tests/mcp/test_search_nodes.py | 39 ++ 16 files changed, 779 insertions(+), 416 deletions(-) create mode 100644 tests/mcp/conftest.py create mode 100644 tests/mcp/test_add_observations.py create mode 100644 tests/mcp/test_create_entitites.py create mode 100644 tests/mcp/test_create_relations.py create mode 100644 tests/mcp/test_delete_entities.py create mode 100644 tests/mcp/test_delete_observations.py create mode 100644 tests/mcp/test_delete_relations.py create mode 100644 tests/mcp/test_list_tools.py create mode 100644 tests/mcp/test_open_nodes.py create mode 100644 tests/mcp/test_search_nodes.py diff --git a/src/basic_memory/api/routers/knowledge.py b/src/basic_memory/api/routers/knowledge.py index 66b5ef9a..d5f4e9c7 100644 --- a/src/basic_memory/api/routers/knowledge.py +++ b/src/basic_memory/api/routers/knowledge.py @@ -112,13 +112,12 @@ async def delete_entity( return DeleteEntityResponse(deleted=deleted) -@router.post( - "/entities/{entity_id:path}/observations/delete", response_model=DeleteObservationsResponse -) +@router.post("/observations/delete", response_model=DeleteObservationsResponse) async def delete_observations( - entity_id: str, data: DeleteObservationsRequest, memory_service: MemoryServiceDep + data: DeleteObservationsRequest, memory_service: MemoryServiceDep ) -> DeleteObservationsResponse: """Delete observations from an entity.""" + entity_id = data.entity_id deleted = await memory_service.delete_observations(entity_id, data.deletions) return DeleteObservationsResponse(deleted=deleted) diff --git a/src/basic_memory/mcp/server.py b/src/basic_memory/mcp/server.py index 6073b8bd..fdd5fb2d 100644 --- a/src/basic_memory/mcp/server.py +++ b/src/basic_memory/mcp/server.py @@ -160,9 +160,11 @@ class MemoryServer(Server): "open_nodes": lambda c, a: c.post("/knowledge/nodes", json=a), "add_observations": lambda c, a: c.post("/knowledge/observations", json=a), "create_relations": lambda c, a: c.post("/knowledge/relations", json=a), - "delete_entities": lambda c, a: c.post(f"/knowledge/entities/{a['names'][0]}"), - "delete_observations": lambda c, a: c.post("/knowledge/observations", json=a), - "delete_relations": lambda c, a: c.post("/knowledge/entities/", json=a), + "delete_entities": lambda c, a: c.post("/knowledge/entities/delete", json=a), + "delete_observations": lambda c, a: c.post( + "/knowledge/observations/delete", json=a + ), + "delete_relations": lambda c, a: c.post("/knowledge/relations/delete", json=a), } # Get handler for tool diff --git a/src/basic_memory/models.py b/src/basic_memory/models.py index 24ad7f35..bc1bbbaf 100644 --- a/src/basic_memory/models.py +++ b/src/basic_memory/models.py @@ -1,15 +1,18 @@ """Database models for basic-memory.""" + from datetime import datetime from typing import List, Optional -from sqlalchemy import String, DateTime, ForeignKey, Text, Integer, text, UniqueConstraint -from sqlalchemy.orm import Mapped, mapped_column, relationship, DeclarativeBase -from sqlalchemy.ext.asyncio import AsyncAttrs -from basic_memory.utils import normalize_entity_id +from sqlalchemy import String, DateTime, ForeignKey, Text, Integer, text, UniqueConstraint +from sqlalchemy.ext.asyncio import AsyncAttrs +from sqlalchemy.orm import Mapped, mapped_column, relationship, DeclarativeBase + +from basic_memory.utils import sanitize_name class Base(AsyncAttrs, DeclarativeBase): """Base class for all models""" + pass @@ -24,54 +27,45 @@ class Entity(Base): - A description (optional) - A list of observations """ + __tablename__ = "entity" - __table_args__ = ( - UniqueConstraint('entity_type', 'name', name='uix_entity_type_name'), - ) + __table_args__ = (UniqueConstraint("entity_type", "name", name="uix_entity_type_name"),) id: Mapped[str] = mapped_column(String, primary_key=True) name: Mapped[str] = mapped_column(String) entity_type: Mapped[str] = mapped_column(String) description: Mapped[Optional[str]] = mapped_column(Text, nullable=True) - created_at: Mapped[datetime] = mapped_column( - DateTime, - server_default=text('CURRENT_TIMESTAMP') - ) + created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP")) updated_at: Mapped[datetime] = mapped_column( - DateTime, - server_default=text('CURRENT_TIMESTAMP'), - onupdate=text('CURRENT_TIMESTAMP') + DateTime, server_default=text("CURRENT_TIMESTAMP"), onupdate=text("CURRENT_TIMESTAMP") ) # Relationships observations: Mapped[List["Observation"]] = relationship( - "Observation", - back_populates="entity", - cascade="all, delete-orphan" + "Observation", back_populates="entity", cascade="all, delete-orphan" ) outgoing_relations: Mapped[List["Relation"]] = relationship( "Relation", foreign_keys="[Relation.from_id]", back_populates="from_entity", - cascade="all, delete-orphan" + cascade="all, delete-orphan", ) incoming_relations: Mapped[List["Relation"]] = relationship( "Relation", foreign_keys="[Relation.to_id]", back_populates="to_entity", - cascade="all, delete-orphan" + cascade="all, delete-orphan", ) @property def relations(self): return self.outgoing_relations + self.incoming_relations - @classmethod def generate_id(cls, entity_type: str, name: str) -> str: """Generate a filesystem path-based ID for this entity.""" # Use common normalization for filesystem safety - safe_name = normalize_entity_id(name) + safe_name = sanitize_name(name) return f"{entity_type}/{safe_name}" def get_file_path(self) -> str: @@ -90,26 +84,19 @@ class Observation(Base): - Can be added or removed independently - Should be atomic (one fact per observation) """ + __tablename__ = "observation" id: Mapped[int] = mapped_column(Integer, primary_key=True) entity_id: Mapped[str] = mapped_column( # Reference Entity.id which is text - String, - ForeignKey("entity.id", ondelete="CASCADE"), - index=True + String, ForeignKey("entity.id", ondelete="CASCADE"), index=True ) content: Mapped[str] = mapped_column(String) - created_at: Mapped[datetime] = mapped_column( - DateTime, - server_default=text('CURRENT_TIMESTAMP') - ) + created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP")) context: Mapped[str | None] = mapped_column(String, nullable=True) # Relationships - entity: Mapped[Entity] = relationship( - "Entity", - back_populates="observations" - ) + entity: Mapped[Entity] = relationship("Entity", back_populates="observations") def __repr__(self) -> str: content = self.content[:50] + "..." if len(self.content) > 50 else self.content @@ -121,37 +108,27 @@ class Relation(Base): Relations define directed connections between entities. They are always stored in active voice and describe how entities interact or relate to each other. """ + __tablename__ = "relation" id: Mapped[int] = mapped_column(Integer, primary_key=True) from_id: Mapped[str] = mapped_column( # Reference Entity.id which is text - String, - ForeignKey("entity.id", ondelete="CASCADE"), - index=True + String, ForeignKey("entity.id", ondelete="CASCADE"), index=True ) to_id: Mapped[str] = mapped_column( # Reference Entity.id which is text - String, - ForeignKey("entity.id", ondelete="CASCADE"), - index=True + String, ForeignKey("entity.id", ondelete="CASCADE"), index=True ) relation_type: Mapped[str] = mapped_column(String) - created_at: Mapped[datetime] = mapped_column( - DateTime, - server_default=text('CURRENT_TIMESTAMP') - ) + created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP")) context: Mapped[str | None] = mapped_column(String, nullable=True) # Relationships from_entity: Mapped[Entity] = relationship( - "Entity", - foreign_keys=[from_id], - back_populates="outgoing_relations" + "Entity", foreign_keys=[from_id], back_populates="outgoing_relations" ) to_entity: Mapped[Entity] = relationship( - "Entity", - foreign_keys=[to_id], - back_populates="incoming_relations" + "Entity", foreign_keys=[to_id], back_populates="incoming_relations" ) def __repr__(self) -> str: - return f"Relation(id={self.id}, from='{self.from_id}', type='{self.relation_type}', to='{self.to_id}')" \ No newline at end of file + return f"Relation(id={self.id}, from='{self.from_id}', type='{self.relation_type}', to='{self.to_id}')" diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index babad5da..c906f234 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -1,17 +1,26 @@ """Core pydantic models for basic-memory entities, observations, and relations.""" -from typing import List, Optional, Annotated, TypeAlias +from typing import List, Optional, Annotated -from annotated_types import Len +from annotated_types import MinLen, MaxLen from pydantic import BaseModel, ConfigDict, BeforeValidator -from basic_memory.utils import normalize_entity_id +from basic_memory.utils import sanitize_name -# Base Models -Observation: TypeAlias = str + +# Strip whitespace +def strip_whitespace(obs: str) -> str: + return obs.strip() + + +Observation = Annotated[str, BeforeValidator(strip_whitespace), MinLen(1), MaxLen(1000)] + +EntityType = Annotated[str, BeforeValidator(strip_whitespace), MinLen(1), MaxLen(20)] + +RelationType = Annotated[str, BeforeValidator(strip_whitespace), MinLen(1), MaxLen(20)] # Custom field types with validation -EntityId = Annotated[str, BeforeValidator(normalize_entity_id)] +EntityId = Annotated[str, BeforeValidator(sanitize_name)] class Relation(BaseModel): @@ -22,7 +31,7 @@ class Relation(BaseModel): from_id: EntityId to_id: EntityId - relation_type: str + relation_type: RelationType context: Optional[str] = None @@ -35,7 +44,7 @@ class Entity(BaseModel): id: Optional[EntityId] = None name: str - entity_type: str + entity_type: EntityType description: Optional[str] = None observations: List[Observation] = [] relations: List[Relation] = [] @@ -60,19 +69,19 @@ class AddObservationsRequest(BaseModel): class CreateEntityRequest(BaseModel): """Request schema for create_entities tool.""" - entities: Annotated[List[Entity], Len(min_length=1)] + entities: Annotated[List[Entity], MinLen(1)] class SearchNodesRequest(BaseModel): """Request schema for search_nodes tool.""" - query: str + query: Annotated[str, MinLen(1), MaxLen(200)] class OpenNodesRequest(BaseModel): """Request schema for open_nodes tool.""" - names: Annotated[List[EntityId], Len(min_length=1)] + names: Annotated[List[EntityId], MinLen(1)] class CreateRelationsRequest(BaseModel): @@ -87,7 +96,7 @@ class CreateRelationsRequest(BaseModel): class DeleteEntityRequest(BaseModel): """Request schema for delete_entities tool.""" - entity_ids: List[EntityId] + entity_ids: Annotated[List[EntityId], MinLen(1)] class DeleteRelationsRequest(BaseModel): @@ -100,7 +109,7 @@ class DeleteObservationsRequest(BaseModel): """Request schema for delete_observations tool.""" entity_id: EntityId - deletions: List[Observation] + deletions: Annotated[List[Observation], MinLen(1)] # response output models @@ -134,7 +143,7 @@ class RelationResponse(Relation, SQLAlchemyModel): class EntityResponse(SQLAlchemyModel): """Schema for entity data returned from the service.""" - id: EntityId + id: str name: str entity_type: str description: Optional[str] = None diff --git a/src/basic_memory/utils.py b/src/basic_memory/utils.py index 8d6f3b32..1988aefe 100644 --- a/src/basic_memory/utils.py +++ b/src/basic_memory/utils.py @@ -1,13 +1,29 @@ """Utility functions for basic-memory.""" -def normalize_entity_id(entity_id: str) -> str: +import re +import unicodedata + + +def sanitize_name(name: str) -> str: """ - Normalize an entity ID by converting to lowercase and replacing spaces with underscores. - - Args: - entity_id: Raw entity ID to normalize - - Returns: - Normalized entity ID suitable for filesystem and database use + Sanitize a name for filesystem use: + - Convert to lowercase + - Replace spaces/punctuation with underscores + - Remove emojis and other special characters + - Collapse multiple underscores + - Trim leading/trailing underscores """ - return entity_id.lower().replace(" ", "_") \ No newline at end of file + # Normalize unicode to compose characters where possible + name = unicodedata.normalize("NFKD", name) + # Remove emojis and other special characters, keep only letters, numbers, spaces + name = "".join(c for c in name if c.isalnum() or c.isspace()) + # Replace spaces with underscores + name = name.replace(" ", "_") + # Remove newline + name = name.replace("\n", "") + # Convert to lowercase + name = name.lower() + # Collapse multiple underscores and trim + name = re.sub(r"_+", "_", name).strip("_") + + return name diff --git a/tests/mcp/conftest.py b/tests/mcp/conftest.py new file mode 100644 index 00000000..dbb94bf1 --- /dev/null +++ b/tests/mcp/conftest.py @@ -0,0 +1,66 @@ +"""Tests for the MCP server implementation using FastAPI TestClient.""" + +import pytest_asyncio +from fastapi import FastAPI +from httpx import AsyncClient, ASGITransport + +from basic_memory.api.app import app as fastapi_app +from basic_memory.deps import get_project_config, get_engine_factory +from basic_memory.mcp.server import MemoryServer + + +@pytest_asyncio.fixture +def app(test_config, engine_session_factory) -> FastAPI: + """Create test FastAPI application.""" + app = fastapi_app + app.dependency_overrides[get_project_config] = lambda: test_config + app.dependency_overrides[get_engine_factory] = lambda: engine_session_factory + return app + + +@pytest_asyncio.fixture() +async def server(app) -> MemoryServer: + server = MemoryServer() + await server.setup() + return server + + +@pytest_asyncio.fixture +async def client(app: FastAPI): + """Create test client that both MCP and tests will use.""" + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + yield client + + +@pytest_asyncio.fixture +def test_entity_data(): + """Sample data for creating a test entity.""" + return { + "entities": [ + { + "name": "Test Entity", + "entity_type": "test", + "description": "", # Empty string instead of None + "observations": ["This is a test observation"], + } + ] + } + + +@pytest_asyncio.fixture +def test_directory_entity_data(): + """Real data that caused failure in the tool.""" + return { + "entities": [ + { + "name": "Directory Organization", + "entity_type": "memory", + "description": "Implemented filesystem organization by entity type", + "observations": [ + "Files are now organized by type using directories like entities/project/basic_memory", + "Entity IDs match filesystem paths for better mental model", + "Fixed path handling bugs by adding consistent get_entity_path helper", + ], + } + ] + } diff --git a/tests/mcp/test_add_observations.py b/tests/mcp/test_add_observations.py new file mode 100644 index 00000000..81123f81 --- /dev/null +++ b/tests/mcp/test_add_observations.py @@ -0,0 +1,39 @@ +"""Tests for the MCP server implementation using FastAPI TestClient.""" + +import pytest +from mcp.types import EmbeddedResource + +from basic_memory.schemas import CreateEntityResponse, AddObservationsResponse + + +@pytest.mark.asyncio +async def test_add_observations(test_entity_data, client, server): + """Test adding observations to an existing entity.""" + + # First create an entity + create_result = await server.handle_call_tool("create_entities", test_entity_data) + create_response = CreateEntityResponse.model_validate_json(create_result[0].resource.text) + entity_id = create_response.entities[0].id + + # Add new observation + result = await server.handle_call_tool( + "add_observations", {"entity_id": entity_id, "observations": ["A new observation"]} + ) + + # Verify response format + assert len(result) == 1 + assert isinstance(result[0], EmbeddedResource) + assert result[0].type == "resource" + + # Verify observation was added + response = AddObservationsResponse.model_validate_json(result[0].resource.text) + assert response.entity_id == entity_id + assert len(response.observations) == 1 + assert response.observations[0].content == "A new observation" + + # Verify through API + api_response = await client.get(f"/knowledge/entities/{entity_id}") + assert api_response.status_code == 200 + entity = api_response.json() + assert len(entity["observations"]) == 2 # Original + new + assert "A new observation" in [o["content"] for o in entity["observations"]] diff --git a/tests/mcp/test_create_entitites.py b/tests/mcp/test_create_entitites.py new file mode 100644 index 00000000..b75fd4ee --- /dev/null +++ b/tests/mcp/test_create_entitites.py @@ -0,0 +1,107 @@ +"""Tests for the MCP server implementation using FastAPI TestClient.""" + +import pytest +from mcp.types import EmbeddedResource + +from basic_memory.mcp.server import MIME_TYPE +from basic_memory.schemas import CreateEntityResponse, SearchNodesResponse + + +@pytest.mark.asyncio +async def test_create_single_entity(server): + """Test creating a single entity.""" + entity_data = { + "entities": [ + {"name": "SingleTest", "entity_type": "test", "observations": ["Test observation"]} + ] + } + + result = await server.handle_call_tool("create_entities", entity_data) + + # Verify response format + assert len(result) == 1 + assert isinstance(result[0], EmbeddedResource) + assert result[0].type == "resource" + assert result[0].resource.mimeType == MIME_TYPE + + # Verify entity creation + response = CreateEntityResponse.model_validate_json(result[0].resource.text) + assert len(response.entities) == 1 + entity = response.entities[0] + assert entity.name == "SingleTest" + assert entity.entity_type == "test" + assert len(entity.observations) == 1 + assert entity.observations[0].content == "Test observation" + assert entity.id == "test/singletest" + + # Verify entity can be found via search + search_result = await server.handle_call_tool("search_nodes", {"query": "SingleTest"}) + search_response = SearchNodesResponse.model_validate_json(search_result[0].resource.text) + assert len(search_response.matches) == 1 + assert search_response.matches[0].name == "SingleTest" + + +@pytest.mark.asyncio +async def test_create_multiple_entities(server): + """Test creating multiple entities in one call.""" + entity_data = { + "entities": [ + {"name": "BulkTest1", "entity_type": "test", "observations": ["First bulk test"]}, + {"name": "BulkTest2", "entity_type": "test", "observations": ["Second bulk test"]}, + {"name": "BulkTest3", "entity_type": "demo", "observations": ["Third bulk test"]}, + ] + } + + result = await server.handle_call_tool("create_entities", entity_data) + + # Verify response + assert len(result) == 1 + response = CreateEntityResponse.model_validate_json(result[0].resource.text) + + # Verify all entities were created + assert len(response.entities) == 3 + + # Check specific entities + entities = {e.name: e for e in response.entities} + assert "BulkTest1" in entities + assert "BulkTest2" in entities + assert "BulkTest3" in entities + + # Verify IDs were generated correctly + assert entities["BulkTest1"].id == "test/bulktest1" + assert entities["BulkTest2"].id == "test/bulktest2" + assert entities["BulkTest3"].id == "demo/bulktest3" + + # Verify observations were saved + assert len(entities["BulkTest1"].observations) == 1 + assert entities["BulkTest1"].observations[0].content == "First bulk test" + + # Verify entities can be found via search + search_result = await server.handle_call_tool("search_nodes", {"query": "BulkTest"}) + search_response = SearchNodesResponse.model_validate_json(search_result[0].resource.text) + assert len(search_response.matches) == 3 + + +@pytest.mark.asyncio +async def test_create_entity_with_all_fields(server): + """Test creating entity with all possible fields populated.""" + entity_data = { + "entities": [ + { + "name": "FullEntity", + "entity_type": "test", + "description": "A complete test entity", + "observations": ["First observation", "Second observation"], + } + ] + } + + result = await server.handle_call_tool("create_entities", entity_data) + response = CreateEntityResponse.model_validate_json(result[0].resource.text) + + entity = response.entities[0] + assert entity.name == "FullEntity" + assert entity.description == "A complete test entity" + assert len(entity.observations) == 2 + assert entity.observations[0].content == "First observation" + assert entity.observations[1].content == "Second observation" diff --git a/tests/mcp/test_create_relations.py b/tests/mcp/test_create_relations.py new file mode 100644 index 00000000..3278e2b7 --- /dev/null +++ b/tests/mcp/test_create_relations.py @@ -0,0 +1,43 @@ +"""Tests for the MCP server implementation using FastAPI TestClient.""" + +import pytest + +from basic_memory.schemas import SearchNodesResponse +from basic_memory.utils import sanitize_name + + +@pytest.mark.asyncio +async def test_create_relations(test_entity_data, client, server): + """Test creating relations between entities.""" + # Create two test entities + entity_data = { + "entities": [ + {"name": "TestEntityA", "entity_type": "test", "observations": ["Entity A"]}, + {"name": "TestEntityB", "entity_type": "test", "observations": ["Entity B"]}, + ] + } + + await server.handle_call_tool("create_entities", entity_data) + + # Create relation between them + relation_data = { + "relations": [ + { + "from_id": "test/TestEntityA", + "to_id": "test/TestEntityB", + "relation_type": "relates_to", + } + ] + } + + result = await server.handle_call_tool("create_relations", relation_data) + + # Verify through search + search_result = await server.handle_call_tool("search_nodes", {"query": "TestEntityA"}) + response = SearchNodesResponse.model_validate_json(search_result[0].resource.text) + + 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].relation_type == "relates_to" diff --git a/tests/mcp/test_delete_entities.py b/tests/mcp/test_delete_entities.py new file mode 100644 index 00000000..08f6442f --- /dev/null +++ b/tests/mcp/test_delete_entities.py @@ -0,0 +1,32 @@ +"""Tests for MCP delete_entities tool.""" + +import pytest + +from basic_memory.schemas import SearchNodesResponse +from basic_memory.utils import sanitize_name + + +@pytest.mark.asyncio +async def test_delete_entities(server): + """Test deleting entities.""" + # Create test entities + entities = { + "entities": [ + {"name": "DeleteTest1", "entity_type": "test", "observations": ["To be deleted 1"]}, + {"name": "DeleteTest2", "entity_type": "test", "observations": ["To be deleted 2"]}, + ] + } + await server.handle_call_tool("create_entities", entities) + + # Delete first entity + await server.handle_call_tool( + "delete_entities", {"entity_ids": [sanitize_name("test/DeleteTest1")]} + ) + + # Verify through search + search_result = await server.handle_call_tool("search_nodes", {"query": "DeleteTest"}) + search_response = SearchNodesResponse.model_validate_json(search_result[0].resource.text) + + # Only second entity should remain + assert len(search_response.matches) == 1 + assert search_response.matches[0].name == "DeleteTest2" diff --git a/tests/mcp/test_delete_observations.py b/tests/mcp/test_delete_observations.py new file mode 100644 index 00000000..2adae147 --- /dev/null +++ b/tests/mcp/test_delete_observations.py @@ -0,0 +1,51 @@ +"""Tests for MCP delete_observations tool.""" + +import pytest + +from basic_memory.schemas import CreateEntityResponse, SearchNodesResponse + + +@pytest.mark.asyncio +async def test_delete_observations(server): + """Test deleting specific observations from an entity.""" + # Create entity with multiple observations + entity_data = { + "entities": [{ + "name": "ObsDeleteTest", + "entity_type": "test", + "observations": [ + "Keep this observation", + "Delete this observation", + "Also keep this" + ] + }] + } + create_result = await server.handle_call_tool("create_entities", entity_data) + create_response = CreateEntityResponse.model_validate_json(create_result[0].resource.text) + entity_id = create_response.entities[0].id + + # Delete specific observation + await server.handle_call_tool( + "delete_observations", + { + "entity_id": entity_id, + "deletions": ["Delete this observation"] + } + ) + + # Verify through search + search_result = await server.handle_call_tool( + "search_nodes", + {"query": "ObsDeleteTest"} + ) + search_response = SearchNodesResponse.model_validate_json( + search_result[0].resource.text + ) + + # Check remaining observations + entity = search_response.matches[0] + observations = [o.content for o in entity.observations] + assert len(observations) == 2 + assert "Delete this observation" not in observations + assert "Keep this observation" in observations + assert "Also keep this" in observations \ No newline at end of file diff --git a/tests/mcp/test_delete_relations.py b/tests/mcp/test_delete_relations.py new file mode 100644 index 00000000..bbaf8c20 --- /dev/null +++ b/tests/mcp/test_delete_relations.py @@ -0,0 +1,53 @@ +"""Tests for MCP delete_relations tool.""" + +import pytest + +from basic_memory.schemas import SearchNodesResponse +from basic_memory.utils import sanitize_name + + +@pytest.mark.asyncio +async def test_delete_relations(server): + """Test deleting relations between entities.""" + # Create test entities with relation + entities = { + "entities": [ + {"name": "RelSource", "entity_type": "test", "observations": ["Source entity"]}, + {"name": "RelTarget", "entity_type": "test", "observations": ["Target entity"]}, + ] + } + await server.handle_call_tool("create_entities", entities) + + # Create relation + relation = { + "relations": [ + { + "from_id": sanitize_name("test/RelSource"), + "to_id": sanitize_name("test/RelTarget"), + "relation_type": "relates_to", + } + ] + } + await server.handle_call_tool("create_relations", relation) + + # Delete the relation + await server.handle_call_tool( + "delete_relations", + { + "relations": [ + { + "from_id": sanitize_name("test/RelSource"), + "to_id": sanitize_name("test/RelTarget"), + "relation_type": "relates_to", + } + ] + }, + ) + + # Verify through search + 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 + assert len(search_response.matches) == 1 + assert len(search_response.matches[0].relations) == 0 diff --git a/tests/mcp/test_list_tools.py b/tests/mcp/test_list_tools.py new file mode 100644 index 00000000..35db2153 --- /dev/null +++ b/tests/mcp/test_list_tools.py @@ -0,0 +1,30 @@ +"""Tests for the MCP server implementation using FastAPI TestClient.""" + +import pytest + + +@pytest.mark.asyncio +async def test_list_tools(server): + """Test that server exposes expected tools.""" + + tools = await server.handle_list_tools() + + # Check each expected tool is present + expected_tools = { + "create_entities", + "search_nodes", + "open_nodes", + "add_observations", + "create_relations", + "delete_entities", + "delete_observations", + "delete_relations", + } + + found_tools = {t.name: t for t in tools} + assert found_tools.keys() == expected_tools + + # Verify schemas include required fields + search_schema = found_tools["search_nodes"].inputSchema + assert "query" in search_schema["properties"] + assert search_schema["required"] == ["query"] diff --git a/tests/mcp/test_mcp_server.py b/tests/mcp/test_mcp_server.py index e29f897d..935daa83 100644 --- a/tests/mcp/test_mcp_server.py +++ b/tests/mcp/test_mcp_server.py @@ -1,303 +1,10 @@ """Tests for the MCP server implementation using FastAPI TestClient.""" import pytest -import pytest_asyncio -from fastapi import FastAPI -from httpx import AsyncClient, ASGITransport from mcp.shared.exceptions import McpError -from mcp.types import EmbeddedResource, INVALID_PARAMS +from mcp.types import INVALID_PARAMS -from basic_memory.api.app import app as fastapi_app -from basic_memory.deps import get_project_config, get_engine_factory -from basic_memory.mcp.server import MemoryServer, MIME_TYPE, BASIC_MEMORY_URI -from basic_memory.schemas import CreateEntityResponse, SearchNodesResponse, AddObservationsResponse -from basic_memory.utils import normalize_entity_id - - -@pytest_asyncio.fixture -def app(test_config, engine_session_factory) -> FastAPI: - """Create test FastAPI application.""" - app = fastapi_app - app.dependency_overrides[get_project_config] = lambda: test_config - app.dependency_overrides[get_engine_factory] = lambda: engine_session_factory - return app - - -@pytest_asyncio.fixture() -async def server(app) -> MemoryServer: - server = MemoryServer() - await server.setup() - return server - - -@pytest_asyncio.fixture -async def client(app: FastAPI): - """Create test client that both MCP and tests will use.""" - async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: - yield client - - -@pytest_asyncio.fixture -def test_entity_data(): - """Sample data for creating a test entity.""" - return { - "entities": [ - { - "name": "Test Entity", - "entity_type": "test", - "description": "", # Empty string instead of None - "observations": ["This is a test observation"], - } - ] - } - - -@pytest_asyncio.fixture -def test_directory_entity_data(): - """Real data that caused failure in the tool.""" - return { - "entities": [ - { - "name": "Directory Organization", - "entity_type": "memory", - "description": "Implemented filesystem organization by entity type", - "observations": [ - "Files are now organized by type using directories like entities/project/basic_memory", - "Entity IDs match filesystem paths for better mental model", - "Fixed path handling bugs by adding consistent get_entity_path helper", - ], - } - ] - } - - -@pytest.mark.asyncio -async def test_create_single_entity(server): - """Test creating a single entity.""" - entity_data = { - "entities": [ - {"name": "SingleTest", "entity_type": "test", "observations": ["Test observation"]} - ] - } - - result = await server.handle_call_tool("create_entities", entity_data) - - # Verify response format - assert len(result) == 1 - assert isinstance(result[0], EmbeddedResource) - assert result[0].type == "resource" - assert result[0].resource.mimeType == MIME_TYPE - - # Verify entity creation - response = CreateEntityResponse.model_validate_json(result[0].resource.text) - assert len(response.entities) == 1 - entity = response.entities[0] - assert entity.name == "SingleTest" - assert entity.entity_type == "test" - assert len(entity.observations) == 1 - assert entity.observations[0].content == "Test observation" - assert entity.id == "test/singletest" - - # Verify entity can be found via search - search_result = await server.handle_call_tool("search_nodes", {"query": "SingleTest"}) - search_response = SearchNodesResponse.model_validate_json(search_result[0].resource.text) - assert len(search_response.matches) == 1 - assert search_response.matches[0].name == "SingleTest" - - -@pytest.mark.asyncio -async def test_create_multiple_entities(server): - """Test creating multiple entities in one call.""" - entity_data = { - "entities": [ - {"name": "BulkTest1", "entity_type": "test", "observations": ["First bulk test"]}, - {"name": "BulkTest2", "entity_type": "test", "observations": ["Second bulk test"]}, - {"name": "BulkTest3", "entity_type": "demo", "observations": ["Third bulk test"]}, - ] - } - - result = await server.handle_call_tool("create_entities", entity_data) - - # Verify response - assert len(result) == 1 - response = CreateEntityResponse.model_validate_json(result[0].resource.text) - - # Verify all entities were created - assert len(response.entities) == 3 - - # Check specific entities - entities = {e.name: e for e in response.entities} - assert "BulkTest1" in entities - assert "BulkTest2" in entities - assert "BulkTest3" in entities - - # Verify IDs were generated correctly - assert entities["BulkTest1"].id == "test/bulktest1" - assert entities["BulkTest2"].id == "test/bulktest2" - assert entities["BulkTest3"].id == "demo/bulktest3" - - # Verify observations were saved - assert len(entities["BulkTest1"].observations) == 1 - assert entities["BulkTest1"].observations[0].content == "First bulk test" - - # Verify entities can be found via search - search_result = await server.handle_call_tool("search_nodes", {"query": "BulkTest"}) - search_response = SearchNodesResponse.model_validate_json(search_result[0].resource.text) - assert len(search_response.matches) == 3 - - -@pytest.mark.asyncio -async def test_create_entity_with_all_fields(server): - """Test creating entity with all possible fields populated.""" - entity_data = { - "entities": [ - { - "name": "FullEntity", - "entity_type": "test", - "description": "A complete test entity", - "observations": ["First observation", "Second observation"], - } - ] - } - - result = await server.handle_call_tool("create_entities", entity_data) - response = CreateEntityResponse.model_validate_json(result[0].resource.text) - - entity = response.entities[0] - assert entity.name == "FullEntity" - assert entity.description == "A complete test entity" - assert len(entity.observations) == 2 - assert entity.observations[0].content == "First observation" - assert entity.observations[1].content == "Second observation" - - -@pytest.mark.asyncio -async def test_list_tools(server): - """Test that server exposes expected tools.""" - - tools = await server.handle_list_tools() - - # Check each expected tool is present - expected_tools = { - "create_entities", - "search_nodes", - "open_nodes", - "add_observations", - "create_relations", - "delete_entities", - "delete_observations", - "delete_relations", - } - - found_tools = {t.name: t for t in tools} - assert found_tools.keys() == expected_tools - - # Verify schemas include required fields - search_schema = found_tools["search_nodes"].inputSchema - assert "query" in search_schema["properties"] - assert search_schema["required"] == ["query"] - - -@pytest.mark.asyncio -async def test_search_nodes(test_entity_data, client, server): - """Test searching for an entity after creating it.""" - - # First create an entity - await server.handle_call_tool("create_entities", test_entity_data) - - # Then search for it - result = await server.handle_call_tool("search_nodes", {"query": "Test Entity"}) - - # Verify response format - assert len(result) == 1 - assert isinstance(result[0], EmbeddedResource) - assert result[0].type == "resource" - 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 - - # Verify search results - response = SearchNodesResponse.model_validate_json(result[0].resource.text) - assert len(response.matches) == 1 - assert response.matches[0].name == "Test Entity" - assert response.query == "Test Entity" - - # Verify through API - api_response = await client.post("/knowledge/search", json={"query": "Test Entity"}) - assert api_response.status_code == 200 - data = api_response.json() - assert len(data["matches"]) == 1 - assert data["matches"][0]["name"] == "Test Entity" - - -@pytest.mark.asyncio -async def test_add_observations(test_entity_data, client, server): - """Test adding observations to an existing entity.""" - - # First create an entity - create_result = await server.handle_call_tool("create_entities", test_entity_data) - create_response = CreateEntityResponse.model_validate_json(create_result[0].resource.text) - entity_id = create_response.entities[0].id - - # Add new observation - result = await server.handle_call_tool( - "add_observations", {"entity_id": entity_id, "observations": ["A new observation"]} - ) - - # Verify response format - assert len(result) == 1 - assert isinstance(result[0], EmbeddedResource) - assert result[0].type == "resource" - - # Verify observation was added - response = AddObservationsResponse.model_validate_json(result[0].resource.text) - assert response.entity_id == entity_id - assert len(response.observations) == 1 - assert response.observations[0].content == "A new observation" - - # Verify through API - api_response = await client.get(f"/knowledge/entities/{entity_id}") - assert api_response.status_code == 200 - entity = api_response.json() - assert len(entity["observations"]) == 2 # Original + new - assert "A new observation" in [o["content"] for o in entity["observations"]] - - -@pytest.mark.asyncio -async def test_create_relations(test_entity_data, client, server): - """Test creating relations between entities.""" - # Create two test entities - entity_data = { - "entities": [ - {"name": "TestEntityA", "entity_type": "test", "observations": ["Entity A"]}, - {"name": "TestEntityB", "entity_type": "test", "observations": ["Entity B"]}, - ] - } - - await server.handle_call_tool("create_entities", entity_data) - - # Create relation between them - relation_data = { - "relations": [ - { - "from_id": "test/TestEntityA", - "to_id": "test/TestEntityB", - "relation_type": "relates_to", - } - ] - } - - result = await server.handle_call_tool("create_relations", relation_data) - - # Verify through search - search_result = await server.handle_call_tool("search_nodes", {"query": "TestEntityA"}) - response = SearchNodesResponse.model_validate_json(search_result[0].resource.text) - - assert len(response.matches) == 1 - entity = response.matches[0] - assert len(entity.relations) == 1 - assert entity.relations[0].to_id == normalize_entity_id("test/TestEntityB") - assert entity.relations[0].relation_type == "relates_to" +from basic_memory.schemas import CreateEntityResponse, SearchNodesResponse @pytest.mark.asyncio @@ -309,56 +16,188 @@ async def test_invalid_tool_name(server): @pytest.mark.asyncio -class TestInputValidation: - """Test input validation for various tools.""" +async def test_missing_required_field(server): + """Test validation when required fields are missing.""" + with pytest.raises(McpError) as exc: + await server.handle_call_tool("search_nodes", {}) + assert "query" in str(exc.value).lower() - async def test_missing_required_field(self, server): - """Test validation when required fields are missing.""" + with pytest.raises(McpError) as exc: + await server.handle_call_tool("create_entities", {}) + assert "entities" in str(exc.value).lower() - with pytest.raises(McpError) as exc: - await server.handle_call_tool("search_nodes", {}) - assert "query" in str(exc.value).lower() - with pytest.raises(McpError) as exc: - await server.handle_call_tool("create_entities", {}) - assert "entities" in str(exc.value).lower() +@pytest.mark.asyncio +async def test_empty_arrays(server): + """Test validation of array fields that can't be empty.""" + with pytest.raises(McpError) as exc: + await server.handle_call_tool("create_entities", {"entities": []}) + assert INVALID_PARAMS == exc.value.args[0] - async def test_empty_arrays(self, server): - """Test validation of array fields that can't be empty.""" + with pytest.raises(McpError) as exc: + await server.handle_call_tool("open_nodes", {"names": []}) + assert INVALID_PARAMS == exc.value.args[0] - with pytest.raises(McpError) as exc: - await server.handle_call_tool("create_entities", {"entities": []}) - assert INVALID_PARAMS == exc.value.args[0] - with pytest.raises(McpError) as exc: - await server.handle_call_tool("open_nodes", {"names": []}) - assert INVALID_PARAMS == exc.value.args[0] +@pytest.mark.asyncio +async def test_invalid_field_types(server): + """Test validation when fields have wrong types.""" + with pytest.raises(McpError) as exc: + await server.handle_call_tool("search_nodes", {"query": 123}) + assert "str" in str(exc.value).lower() - async def test_invalid_field_types(self, server): - """Test validation when fields have wrong types.""" + with pytest.raises(McpError) as exc: + await server.handle_call_tool("create_entities", {"entities": "not an array"}) + assert "array" in str(exc.value).lower() or "list" in str(exc.value).lower() - with pytest.raises(McpError) as exc: - await server.handle_call_tool("search_nodes", {"query": 123}) - assert "str" in str(exc.value).lower() - with pytest.raises(McpError) as exc: - await server.handle_call_tool("create_entities", {"entities": "not an array"}) - assert "array" in str(exc.value).lower() or "list" in str(exc.value).lower() +@pytest.mark.asyncio +async def test_invalid_nested_fields(server): + """Test validation of nested object fields.""" + with pytest.raises(McpError) as exc: + await server.handle_call_tool( + "create_entities", + { + "entities": [ + { + "name": "Test", + # Missing required entity_type + "observations": [], + } + ] + }, + ) + assert "entity_type" in str(exc.value).lower() - async def test_invalid_nested_fields(self, server): - """Test validation of nested object fields.""" - with pytest.raises(McpError) as exc: - await server.handle_call_tool( - "create_entities", - { - "entities": [ - { - "name": "Test", - # Missing required entity_type - "observations": [], - } - ] - }, - ) - assert "entity_type" in str(exc.value).lower() +@pytest.mark.asyncio +async def test_invalid_relation_format_to_id(server): + """Test validation of relation data.""" + with pytest.raises(McpError) as exc: + await server.handle_call_tool( + "create_relations", + { + "relations": [ + { + "from_id": "test/entity1", + # Missing to_id + "relation_type": "relates_to", + } + ] + }, + ) + assert "to_id" in str(exc.value).lower() + + +@pytest.mark.asyncio +async def test_invalid_relation_format_relation_type(server): + # Invalid relation type + with pytest.raises(McpError) as exc: + await server.handle_call_tool( + "create_relations", + { + "relations": [ + { + "from_id": "test/entity1", + "to_id": "test/entity2", + "relation_type": "", # Empty relation type + } + ] + }, + ) + assert "relation_type" in str(exc.value).lower() + + +@pytest.mark.asyncio +async def test_observation_validation_len(server): + """Test validation specific to observations.""" + # Empty observations + with pytest.raises(McpError) as exc: + await server.handle_call_tool( + "add_observations", + { + "entity_id": "test/entity1", + "observations": ["", ""], # Empty observations + }, + ) + assert "observations" in str(exc.value).lower() + + +@pytest.mark.asyncio +async def test_observation_validation_delete(server): + # Empty deletions + with pytest.raises(McpError) as exc: + await server.handle_call_tool( + "delete_observations", + { + "entity_id": "test/entity1", + "deletions": [], # Empty deletions + }, + ) + assert INVALID_PARAMS == exc.value.args[0] + + +@pytest.mark.asyncio +async def test_edge_case_validation_search_len(server): + """Test edge cases in validation.""" + # Very long strings + with pytest.raises(McpError) as exc: + await server.handle_call_tool( + "search_nodes", + {"query": "x" * 10000}, # Extremely long query + ) + + +@pytest.mark.asyncio +async def test_edge_case_validation_name_sanitization(server): + """Test that entity names are properly sanitized for IDs.""" + # Test cases for different sanitization scenarios + test_cases = [ + { + "name": "🧪 FOO & File (1)", # Emoji and special chars + "expected_id": "test/foo_file_1", + }, + { + "name": "BARR Multiple Spaces", # Multiple spaces + "expected_id": "test/barr_multiple_spaces", + }, + { + "name": "LOTSOF@#$Special&*Chars", # Special characters + "expected_id": "test/lotsofspecialchars", + }, + { + "name": "\x00null", # Null byte + "expected_id": "test/null", + }, + { + "name": "\nline", # Newline + "expected_id": "test/line", + }, + ] + + for test in test_cases: + result = await server.handle_call_tool( + "create_entities", + { + "entities": [ + { + "name": test["name"], + "entity_type": "test", + "observations": [], + } + ] + }, + ) + + response = CreateEntityResponse.model_validate_json(result[0].resource.text) + entity = response.entities[0] + + # Original name should be preserved + assert entity.name == test["name"] + # ID should be sanitized + assert entity.id == test["expected_id"] + + # Verify we can find it with original name + search_result = await server.handle_call_tool("search_nodes", {"query": test["name"]}) + search_response = SearchNodesResponse.model_validate_json(search_result[0].resource.text) + assert len(search_response.matches) > 0 diff --git a/tests/mcp/test_open_nodes.py b/tests/mcp/test_open_nodes.py new file mode 100644 index 00000000..30ea4ec7 --- /dev/null +++ b/tests/mcp/test_open_nodes.py @@ -0,0 +1,61 @@ +"""Tests for MCP open_nodes tool.""" + +import pytest +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 +async def test_open_nodes(server): + """Test retrieving specific nodes by name.""" + # Create test entities + entity_data = { + "entities": [ + { + "name": "OpenTestA", + "entity_type": "test", + "observations": ["First test entity"], + }, + { + "name": "OpenTestB", + "entity_type": "test", + "observations": ["Second test entity"], + }, + { + "name": "OpenTestC", + "entity_type": "test", + "observations": ["Third test entity"], + }, + ] + } + + await server.handle_call_tool("create_entities", entity_data) + + # Open specific nodes + result = await server.handle_call_tool( + "open_nodes", {"names": ["test/opentesta", "test/opentestb"]} + ) + + # Verify response format + assert len(result) == 1 + assert isinstance(result[0], EmbeddedResource) + assert result[0].type == "resource" + assert result[0].resource.mimeType == MIME_TYPE + + # Verify entities returned + response = OpenNodesResponse.model_validate_json(result[0].resource.text) + assert len(response.entities) == 2 + + # Entities should be returned in same order as requested + assert response.entities[0].name == "OpenTestA" + assert response.entities[1].name == "OpenTestB" + + # Verify entity content + entity = response.entities[0] + assert entity.id == sanitize_name("test/OpenTestA") + assert entity.entity_type == "test" + assert len(entity.observations) == 1 + assert entity.observations[0].content == "First test entity" diff --git a/tests/mcp/test_search_nodes.py b/tests/mcp/test_search_nodes.py new file mode 100644 index 00000000..159d6140 --- /dev/null +++ b/tests/mcp/test_search_nodes.py @@ -0,0 +1,39 @@ +"""Tests for the MCP server implementation using FastAPI TestClient.""" + +import pytest +from mcp.types import EmbeddedResource + +from basic_memory.mcp.server import MIME_TYPE, BASIC_MEMORY_URI +from basic_memory.schemas import SearchNodesResponse + + +@pytest.mark.asyncio +async def test_search_nodes(test_entity_data, client, server): + """Test searching for an entity after creating it.""" + + # First create an entity + await server.handle_call_tool("create_entities", test_entity_data) + + # Then search for it + result = await server.handle_call_tool("search_nodes", {"query": "Test Entity"}) + + # Verify response format + assert len(result) == 1 + assert isinstance(result[0], EmbeddedResource) + assert result[0].type == "resource" + 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 + + # Verify search results + response = SearchNodesResponse.model_validate_json(result[0].resource.text) + assert len(response.matches) == 1 + assert response.matches[0].name == "Test Entity" + assert response.query == "Test Entity" + + # Verify through API + api_response = await client.post("/knowledge/search", json={"query": "Test Entity"}) + assert api_response.status_code == 200 + data = api_response.json() + assert len(data["matches"]) == 1 + assert data["matches"][0]["name"] == "Test Entity"