mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
MemoryServer wip
This commit is contained in:
phernandez
+188
-64
@@ -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."""
|
||||
|
||||
+21
-13
@@ -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
|
||||
Reference in New Issue
Block a user