admin/basicmachines-co-basic-memory
1
0
Fork 0
mirror of https://github.com/basicmachines-co/basic-memory synced 2026-06-21 13:47:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

MemoryServer wip

Browse Source
This commit is contained in:
phernandez
2024-12-08 14:14:24 -06:00
parent fb21192119
commit e0361d3422
3 changed files with 287 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/basic_memory/mcp/server.py
+188 -64
View File
@@ -1,77 +1,201 @@
"""MCP server implementation for basic-memory."""
from pathlib import Path
from typing import Annotated, List, Dict, Any
from typing import List, Dict, Any
from mcp.server import Server
from mcp.types import Tool, TextContent, METHOD_NOT_FOUND, INVALID_PARAMS, INTERNAL_ERROR
from mcp.shared.exceptions import McpError
from pydantic import BaseModel, Field
from basic_memory import deps
from basic_memory.schemas import EntityIn, ObservationIn, RelationIn
class MemoryServer(Server):
"""Extended server class that exposes handlers for testing."""
def __init__(self):
super().__init__("basic-memory")
self.register_handlers()
def register_handlers(self):
"""Register all handlers with proper decorators."""
@self.list_tools()
async def handle_list_tools() -> List[Tool]:
"""Define the available tools."""
return [
Tool(
name="create_entities",
description="Create multiple new entities in the knowledge graph",
inputSchema={
"type": "object",
"properties": {
"entities": {
"type": "array",
"items": EntityIn.model_json_schema(),
"minItems": 1
}
},
"required": ["entities"]
}
),
Tool(
name="search_nodes",
description="Search for nodes in the knowledge graph",
inputSchema={
"type": "object",
"properties": {
"query": {
"type": "string",
"minLength": 1
}
},
"required": ["query"]
}
),
Tool(
name="open_nodes",
description="Open specific nodes by their names",
inputSchema={
"type": "object",
"properties": {
"names": {
"type": "array",
"items": {"type": "string"},
"minItems": 1
}
},
"required": ["names"]
}
),
Tool(
name="add_observations",
description="Add observations to existing entities",
inputSchema={
"type": "object",
"properties": {
"entityId": {"type": "string"},
"observations": {
"type": "array",
"items": ObservationIn.model_json_schema(),
"minItems": 1
}
},
"required": ["entityId", "observations"]
}
),
Tool(
name="create_relations",
description="Create relations between entities",
inputSchema={
"type": "object",
"properties": {
"relations": {
"type": "array",
"items": RelationIn.model_json_schema(),
"minItems": 1
}
},
"required": ["relations"]
}
),
Tool(
name="delete_entities",
description="Delete entities from the knowledge graph",
inputSchema={
"type": "object",
"properties": {
"names": {
"type": "array",
"items": {"type": "string"},
"minItems": 1
}
},
"required": ["names"]
}
),
Tool(
name="delete_observations",
description="Delete observations from entities",
inputSchema={
"type": "object",
"properties": {
"deletions": {
"type": "array",
"items": {
"type": "object",
"properties": {
"entityName": {"type": "string"},
"observations": {
"type": "array",
"items": {"type": "string"},
"minItems": 1
}
},
"required": ["entityName", "observations"]
}
}
},
"required": ["deletions"]
}
)
]
@self.call_tool()
async def handle_call_tool(name: str, arguments: Dict[str, Any]) -> List[TextContent]:
"""Handle tool calls by delegating to the memory service."""
try:
# Get project path (could come from context in future)
project_path = Path.home() / ".basic-memory" / "projects" / "default"
# Get services with proper lifecycle management
async with deps.get_project_services(project_path) as memory_service:
match name:
case "create_entities":
# Each entity in arguments["entities"] will be validated by EntityIn
result = await memory_service.create_entities(arguments["entities"])
return [TextContent(type="text", text=str(result))]
case "search_nodes":
result = await memory_service.search_nodes(arguments["query"])
return [TextContent(type="text", text=str(result))]
case "open_nodes":
result = await memory_service.open_nodes(arguments["names"])
return [TextContent(type="text", text=str(result))]
case "add_observations":
result = await memory_service.add_observations(arguments)
return [TextContent(type="text", text=str(result))]
case "create_relations":
result = await memory_service.create_relations(arguments["relations"])
return [TextContent(type="text", text=str(result))]
case "delete_entities":
await memory_service.delete_entities(arguments["names"])
return [TextContent(type="text", text="Entities deleted")]
case "delete_observations":
await memory_service.delete_observations(arguments["deletions"])
return [TextContent(type="text", text="Observations deleted")]
case _:
raise McpError(
METHOD_NOT_FOUND,
f"Unknown tool: {name}"
)
except ValueError as e:
raise McpError(INVALID_PARAMS, str(e))
except Exception as e:
raise McpError(INTERNAL_ERROR, str(e))
# Store handlers as instance attributes for testing
self.handle_list_tools = handle_list_tools
self.handle_call_tool = handle_call_tool
# Create server instance
server = Server("basic-memory")
# Parameter models
class CreateEntitiesParams(BaseModel):
"""Parameters for creating entities."""
entities: Annotated[List[Dict[str, Any]], Field(
description="List of entities to create",
min_items=1
)]
class SearchNodesParams(BaseModel):
"""Parameters for searching nodes."""
query: Annotated[str, Field(
description="Search query to match against entities"
)]
@server.list_tools()
async def list_tools() -> List[Tool]:
"""Define the available tools."""
return [
Tool(
name="create_entities",
description="Create multiple new entities in the knowledge graph",
inputSchema=CreateEntitiesParams.model_json_schema()
),
Tool(
name="search_nodes",
description="Search for nodes in the knowledge graph",
inputSchema=SearchNodesParams.model_json_schema()
)
]
@server.call_tool()
async def call_tool(name: str, arguments: Dict[str, Any]) -> List[TextContent]:
"""Handle tool calls by delegating to the memory service."""
try:
# Get project path (could come from context in future)
project_path = Path.home() / ".basic-memory" / "projects" / "default"
# Get services with proper lifecycle management
async with deps.get_project_services(project_path) as memory_service:
match name:
case "create_entities":
params = CreateEntitiesParams(**arguments)
result = await memory_service.create_entities(params.entities)
return [TextContent(type="text", text=str(result))]
case "search_nodes":
params = SearchNodesParams(**arguments)
result = await memory_service.search_nodes(params.query)
return [TextContent(type="text", text=str(result))]
case _:
raise McpError(
METHOD_NOT_FOUND,
f"Unknown tool: {name}"
)
except ValueError as e:
raise McpError(INVALID_PARAMS, str(e))
except Exception as e:
raise McpError(INTERNAL_ERROR, str(e))
server = MemoryServer()
async def run_server():
"""Run the MCP server."""
src/basic_memory/schemas.py
+21 -13
View File
@@ -7,8 +7,7 @@ from datetime import datetime, UTC
from typing import List, Optional, Dict, Any
from uuid import uuid4
from pydantic import BaseModel, model_validator
from pydantic import BaseModel, Field, model_validator
class ObservationIn(BaseModel):
"""Schema for creating a single observation."""
@@ -17,29 +16,37 @@ class ObservationIn(BaseModel):
class ObservationsIn(BaseModel):
"""Schema for adding observations to an entity."""
entity_id: str # Maps to Entity.id
entity_id: str = Field(alias="entityId") # Maps to Entity.id
observations: List[ObservationIn]
class Config:
populate_by_name = True
class ObservationOut(ObservationIn):
"""Schema for observation data returned from the service."""
id: int
class ObservationsOut(BaseModel):
"""Schema for bulk observation operation results."""
entity_id: str
entity_id: str = Field(alias="entityId")
observations: List[ObservationOut]
class Config:
populate_by_name = True
class RelationIn(BaseModel):
"""
Represents a directed edge between entities in the knowledge graph.
Relations are always stored in active voice (e.g. "created", "teaches", etc.)
"""
from_id: str
to_id: str
relation_type: str
from_id: str = Field(alias="fromId")
to_id: str = Field(alias="toId")
relation_type: str = Field(alias="relationType")
context: Optional[str] = None
class Config:
populate_by_name = True
class RelationOut(BaseModel):
id: int
@@ -47,7 +54,7 @@ class EntityBase(BaseModel):
# id assigned at creation via model_validator
id: str
name: str
entity_type: str
entity_type: str = Field(alias="entityType")
@model_validator(mode='before')
@classmethod
@@ -72,12 +79,13 @@ class EntityIn(EntityBase):
observations: List[ObservationIn] = []
relations: List[RelationIn] = []
class Config:
populate_by_name = True
class EntityOut(EntityBase):
"""
Represents a node in our knowledge graph - could be a person, project,
concept, etc. Each entity has a unique name, a type, and a list of
associated observations.
"""
"""Schema for entity data returned from the service."""
observations: List[ObservationOut] = []
relations: List[RelationOut] = []
class Config:
populate_by_name = True
tests/test_mcp_server.py
+78 -30
View File
@@ -1,8 +1,7 @@
"""Tests for the MCP server implementation."""
import pytest
from pathlib import Path
from mcp.types import Tool, TextContent
from mcp.types import TextContent
from basic_memory.mcp import server
@pytest.fixture
@@ -11,49 +10,87 @@ def anyio_backend():
@pytest.fixture
def test_entity_data():
"""Sample data for creating a test entity using camelCase (like MCP will)."""
return {
"entities": [{
"name": "Test Entity",
"name": "Test Entity CamelCase",
"entityType": "test",
"observations": ["This is a test observation"]
"observations": [{"content": "This is a test observation"}]
}]
}
@pytest.fixture
def test_entity_snake_case():
"""Same test data but using snake_case to test schema flexibility."""
return {
"entities": [{
"name": "Test Entity SnakeCase",
"entity_type": "test",
"observations": [{"content": "This is a test observation"}]
}]
}
@pytest.mark.anyio
async def test_list_tools():
"""Test that server exposes expected tools."""
tools = server.list_tools()
tools = await server.handle_list_tools()
assert len(tools) == 2 # We have create_entities and search_nodes for now
# Check each expected tool is present
expected_tools = {
"create_entities", "search_nodes", "open_nodes",
"add_observations", "create_relations",
"delete_entities", "delete_observations"
}
# Verify create_entities tool
create_tool = next(t for t in tools if t.name == "create_entities")
assert create_tool.inputSchema["required"] == ["entities"]
assert "entities" in create_tool.inputSchema["properties"]
found_tools = {t.name: t for t in tools}
assert found_tools.keys() == expected_tools
# Verify search_nodes tool
search_tool = next(t for t in tools if t.name == "search_nodes")
assert search_tool.inputSchema["required"] == ["query"]
assert "query" in search_tool.inputSchema["properties"]
# Verify schemas include required fields
assert "entities" in found_tools["create_entities"].inputSchema["required"]
assert "query" in found_tools["search_nodes"].inputSchema["required"]
@pytest.mark.anyio
async def test_call_create_entities(test_entity_data):
"""Test creating an entity through the tool interface."""
result = await server.call_tool("create_entities", test_entity_data)
async def test_create_entities_camel_case(test_entity_data):
"""Test creating an entity with camelCase data (like from MCP)."""
result = await server.handle_call_tool("create_entities", test_entity_data)
assert len(result) == 1
assert isinstance(result[0], TextContent)
text_content = result[0].text
assert "Test Entity" in text_content
assert "Test Entity" in result[0].text
@pytest.mark.anyio
async def test_call_search_nodes(test_entity_data):
async def test_create_entities_snake_case(test_entity_snake_case):
"""Test creating an entity with snake_case data (like internal usage)."""
result = await server.handle_call_tool("create_entities", test_entity_snake_case)
assert len(result) == 1
assert isinstance(result[0], TextContent)
assert "Test Entity" in result[0].text
@pytest.mark.anyio
async def test_search_nodes(test_entity_data):
"""Test searching for an entity after creating it."""
# First create an entity
await server.call_tool("create_entities", test_entity_data)
await server.handle_call_tool("create_entities", test_entity_data)
# Then search for it
result = await server.call_tool("search_nodes", {"query": "Test Entity"})
result = await server.handle_call_tool("search_nodes", {"query": "Test Entity"})
assert len(result) == 1
assert isinstance(result[0], TextContent)
assert "Test Entity" in result[0].text
@pytest.mark.anyio
async def test_add_observations(test_entity_data):
"""Test adding observations to an existing entity."""
# First create an entity
await server.handle_call_tool("create_entities", test_entity_data)
# Add new observations using camelCase
result = await server.handle_call_tool("add_observations", {
"entityId": "Test Entity",
"observations": [{"content": "A new observation"}]
})
assert len(result) == 1
assert isinstance(result[0], TextContent)
@@ -62,19 +99,30 @@ async def test_call_search_nodes(test_entity_data):
@pytest.mark.anyio
async def test_invalid_tool_name():
"""Test calling a non-existent tool."""
with pytest.raises(Exception) as exc: # We could be more specific about error type
await server.call_tool("not_a_tool", {})
with pytest.raises(Exception) as exc:
await server.handle_call_tool("not_a_tool", {})
assert "Unknown tool" in str(exc.value)
@pytest.mark.anyio
async def test_invalid_parameters():
"""Test calling tools with invalid parameters."""
# Test missing required parameter
"""Test validation with invalid parameters."""
# Test missing required field
with pytest.raises(Exception) as exc:
await server.call_tool("search_nodes", {})
assert "query" in str(exc.value)
await server.handle_call_tool("search_nodes", {})
assert "query" in str(exc.value).lower()
# Test empty entities list
with pytest.raises(Exception) as exc:
await server.call_tool("create_entities", {"entities": []})
assert "min_items" in str(exc.value)
await server.handle_call_tool("create_entities", {"entities": []})
assert "length" in str(exc.value).lower()
# Test invalid case mixing (should pass due to aliases)
result = await server.handle_call_tool("create_entities", {
"entities": [{
"name": "Mixed Case Test",
"entityType": "test", # camelCase
"observations": [{"content": "Testing mixed case"}]
}]
})
assert len(result) == 1
assert "Mixed Case Test" in result[0].text