mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
refactor mcp server to use rest client
This commit is contained in:
phernandez
+74
-286
@@ -1,355 +1,143 @@
|
||||
"""MCP server implementation for basic-memory."""
|
||||
import sys
|
||||
"""MCP server implementation using FastAPI TestClient."""
|
||||
import json
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Any, Optional, Literal, Callable, Awaitable
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
from typing_extensions import TypeAlias
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from httpx import AsyncClient, ASGITransport
|
||||
from mcp import McpError
|
||||
from mcp.server import Server
|
||||
from mcp.types import Tool, EmbeddedResource, TextResourceContents, METHOD_NOT_FOUND, INVALID_PARAMS, INTERNAL_ERROR
|
||||
from mcp.shared.exceptions import McpError
|
||||
from pydantic.networks import AnyUrl
|
||||
from pydantic import TypeAdapter, BaseModel
|
||||
from mcp.types import Tool, EmbeddedResource, TextResourceContents
|
||||
from pydantic import TypeAdapter, AnyUrl
|
||||
|
||||
from basic_memory import db
|
||||
from basic_memory.config import ProjectConfig
|
||||
from basic_memory.fileio import EntityNotFoundError
|
||||
from basic_memory.repository.entity_repository import EntityRepository
|
||||
from basic_memory.repository.observation_repository import ObservationRepository
|
||||
from basic_memory.repository.relation_repository import RelationRepository
|
||||
from basic_memory.api.app import app as fastapi_app
|
||||
from basic_memory.schemas import (
|
||||
# Tool inputs
|
||||
CreateEntityRequest, SearchNodesRequest, OpenNodesRequest,
|
||||
CreateRelationsRequest, DeleteEntityRequest,
|
||||
DeleteObservationsRequest,
|
||||
# Tool responses
|
||||
CreateEntityResponse, SearchNodesResponse, OpenNodesResponse,
|
||||
AddObservationsResponse, CreateRelationsResponse, DeleteEntityResponse,
|
||||
EntityResponse, ObservationResponse, Relation, AddObservationsRequest
|
||||
AddObservationsRequest, CreateRelationsRequest, DeleteEntityRequest,
|
||||
DeleteObservationsRequest
|
||||
)
|
||||
from basic_memory.services import EntityService, ObservationService, RelationService
|
||||
from basic_memory.services.memory_service import MemoryService
|
||||
from loguru import logger
|
||||
|
||||
MIME_TYPE = "application/vnd.basic-memory+json"
|
||||
# URI constants
|
||||
url_validator = TypeAdapter(AnyUrl)
|
||||
BASIC_MEMORY_URI = url_validator.validate_python("basic-memory://response")
|
||||
MIME_TYPE = "application/vnd.basic-memory+json"
|
||||
|
||||
# Define tool name type and handler type
|
||||
ToolName = Literal[
|
||||
"create_entities",
|
||||
"search_nodes",
|
||||
"open_nodes",
|
||||
"add_observations",
|
||||
"create_relations",
|
||||
"delete_entities",
|
||||
"delete_observations"
|
||||
]
|
||||
|
||||
ToolHandler: TypeAlias = Callable[[MemoryService, Dict[str, Any]], Awaitable[EmbeddedResource]]
|
||||
|
||||
@asynccontextmanager
|
||||
async def get_memory_service_session(engine: AsyncEngine, project_path: Path):
|
||||
"""Get all services with proper session and lifecycle management."""
|
||||
async with db.session(engine) as session:
|
||||
# Create repos
|
||||
entity_repo = EntityRepository(session)
|
||||
observation_repo = ObservationRepository(session)
|
||||
relation_repo = RelationRepository(session)
|
||||
|
||||
# Create services
|
||||
entity_service = EntityService(project_path, entity_repo)
|
||||
observation_service = ObservationService(project_path, observation_repo)
|
||||
relation_service = RelationService(project_path, relation_repo)
|
||||
|
||||
# Create memory service
|
||||
memory_service = MemoryService(
|
||||
project_path=project_path,
|
||||
entity_service=entity_service,
|
||||
relation_service=relation_service,
|
||||
observation_service=observation_service
|
||||
)
|
||||
|
||||
yield memory_service
|
||||
|
||||
@asynccontextmanager
|
||||
async def get_project_services(project_path: Path):
|
||||
"""Get all services for a project with full lifecycle management."""
|
||||
async with db.engine(project_path=project_path) as engine:
|
||||
async with get_memory_service_session(engine, project_path) as services:
|
||||
yield services
|
||||
|
||||
def create_response(response: BaseModel) -> EmbeddedResource:
|
||||
"""Create standard MCP response from any response model."""
|
||||
logger.debug(f"Creating MCP response from {response.__class__.__name__}")
|
||||
result = EmbeddedResource(
|
||||
def create_response(data: Dict[str, Any]) -> EmbeddedResource:
|
||||
"""Create standard MCP response."""
|
||||
return EmbeddedResource(
|
||||
type="resource",
|
||||
resource=TextResourceContents(
|
||||
uri=BASIC_MEMORY_URI,
|
||||
mimeType=MIME_TYPE,
|
||||
text=response.model_dump_json()
|
||||
text=json.dumps(data),
|
||||
)
|
||||
)
|
||||
logger.debug(f"Created response: {result}")
|
||||
return result
|
||||
|
||||
|
||||
async def handle_create_entities(
|
||||
service: MemoryService,
|
||||
args: Dict[str, Any]
|
||||
) -> EmbeddedResource:
|
||||
"""Handle create_entities tool call."""
|
||||
# Validate input
|
||||
logger.debug(f"Creating entities with args: {args}")
|
||||
input_args = CreateEntityRequest.model_validate(args)
|
||||
logger.debug(f"Validated input: {len(input_args.entities)} entities")
|
||||
|
||||
# Call service with validated data
|
||||
entities = await service.create_entities(input_args.entities)
|
||||
logger.debug(f"Created {len(entities)} entities")
|
||||
|
||||
# Format response
|
||||
response = CreateEntityResponse(entities=[EntityResponse.model_validate(entity) for entity in entities])
|
||||
logger.debug("Formatted create_entities response")
|
||||
return create_response(response)
|
||||
|
||||
|
||||
async def handle_search_nodes(
|
||||
service: MemoryService,
|
||||
args: Dict[str, Any]
|
||||
) -> EmbeddedResource:
|
||||
"""Handle search_nodes tool call."""
|
||||
logger.debug(f"Searching nodes with query: {args.get('query')}")
|
||||
input_args = SearchNodesRequest.model_validate(args)
|
||||
results = await service.search_nodes(input_args.query)
|
||||
logger.debug(f"Found {len(results)} matches for query '{input_args.query}'")
|
||||
response = SearchNodesResponse(
|
||||
matches=[EntityResponse.model_validate(entity) for entity in results],
|
||||
query=input_args.query
|
||||
)
|
||||
return create_response(response)
|
||||
|
||||
|
||||
async def handle_open_nodes(
|
||||
service: MemoryService,
|
||||
args: Dict[str, Any]
|
||||
) -> EmbeddedResource:
|
||||
"""Handle open_nodes tool call."""
|
||||
logger.debug(f"Opening nodes: {args.get('names')}")
|
||||
input_args = OpenNodesRequest.model_validate(args)
|
||||
entities = await service.open_nodes(input_args.names)
|
||||
logger.debug(f"Opened {len(entities)} entities")
|
||||
response = OpenNodesResponse(entities=[EntityResponse.model_validate(entity) for entity in entities])
|
||||
return create_response(response)
|
||||
|
||||
|
||||
async def handle_add_observations(
|
||||
service: MemoryService,
|
||||
args: Dict[str, Any]
|
||||
) -> EmbeddedResource:
|
||||
"""Handle add_observations tool call."""
|
||||
# Validate input
|
||||
logger.debug(f"Adding observations: {args}")
|
||||
input_args = AddObservationsRequest.model_validate(args)
|
||||
logger.debug(f"Adding {len(input_args.observations)} observations to entity {input_args.entity_id}")
|
||||
|
||||
# Call service with validated data
|
||||
observations = await service.add_observations(input_args)
|
||||
logger.debug(f"Added {len(observations)} observations")
|
||||
|
||||
# Format response
|
||||
response = AddObservationsResponse(
|
||||
entity_id=input_args.entity_id,
|
||||
observations=[ObservationResponse.model_validate(obs) for obs in observations]
|
||||
)
|
||||
return create_response(response)
|
||||
|
||||
|
||||
async def handle_create_relations(
|
||||
service: MemoryService,
|
||||
args: Dict[str, Any]
|
||||
) -> EmbeddedResource:
|
||||
"""Handle create_relations tool call."""
|
||||
# Validate input
|
||||
logger.debug(f"Creating relations: {args}")
|
||||
input_args = CreateRelationsRequest.model_validate(args)
|
||||
logger.debug(f"Creating {len(input_args.relations)} relations")
|
||||
|
||||
# Call service with validated data
|
||||
created = await service.create_relations(input_args.relations)
|
||||
logger.debug(f"Created {len(created)} relations")
|
||||
|
||||
# Format response
|
||||
response = CreateRelationsResponse(relations=[Relation.model_validate(relation) for relation in created])
|
||||
return create_response(response)
|
||||
|
||||
|
||||
async def handle_delete_entities(
|
||||
service: MemoryService,
|
||||
args: Dict[str, Any]
|
||||
) -> EmbeddedResource:
|
||||
"""Handle delete_entities tool call."""
|
||||
logger.debug(f"Deleting entities: {args}")
|
||||
input_args = DeleteEntityRequest.model_validate(args)
|
||||
deleted = await service.delete_entities(input_args.names)
|
||||
logger.debug(f"Deleted entities: {deleted}")
|
||||
response = DeleteEntityResponse(deleted=deleted)
|
||||
return create_response(response)
|
||||
|
||||
|
||||
async def handle_delete_observations(
|
||||
service: MemoryService,
|
||||
args: Dict[str, Any]
|
||||
) -> EmbeddedResource:
|
||||
"""Handle delete_observations tool call."""
|
||||
logger.debug(f"Deleting observations: {args}")
|
||||
return EmbeddedResource()
|
||||
|
||||
|
||||
# Map tool names to handlers
|
||||
TOOL_HANDLERS: Dict[ToolName, ToolHandler] = {
|
||||
"create_entities": handle_create_entities,
|
||||
"search_nodes": handle_search_nodes,
|
||||
"open_nodes": handle_open_nodes,
|
||||
"add_observations": handle_add_observations,
|
||||
"create_relations": handle_create_relations,
|
||||
"delete_entities": handle_delete_entities,
|
||||
"delete_observations": handle_delete_observations,
|
||||
}
|
||||
@asynccontextmanager
|
||||
async def get_client():
|
||||
"""Get FastAPI test client."""
|
||||
async with AsyncClient(
|
||||
transport=ASGITransport(app=fastapi_app),
|
||||
base_url="http://test"
|
||||
) as client:
|
||||
yield client
|
||||
|
||||
|
||||
class MemoryServer(Server):
|
||||
"""Extended server class that exposes handlers for testing."""
|
||||
|
||||
def __init__(self, config: Optional[ProjectConfig] = None):
|
||||
"""MCP server that forwards requests to FastAPI."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__("basic-memory")
|
||||
self.config = config or ProjectConfig()
|
||||
logger.debug(f"Initialized MemoryServer with config: {self.config}")
|
||||
self.register_handlers()
|
||||
|
||||
|
||||
def register_handlers(self):
|
||||
"""Register all handlers with proper decorators."""
|
||||
|
||||
"""Register all handlers."""
|
||||
|
||||
@self.list_tools()
|
||||
async def handle_list_tools() -> List[Tool]:
|
||||
"""Define the available tools."""
|
||||
logger.debug("Listing available tools")
|
||||
tools = [
|
||||
return [
|
||||
Tool(
|
||||
name="create_entities",
|
||||
description="Create multiple new entities in the knowledge graph",
|
||||
description="Create multiple new entities",
|
||||
inputSchema=CreateEntityRequest.model_json_schema()
|
||||
),
|
||||
Tool(
|
||||
name="search_nodes",
|
||||
description="Search for nodes in the knowledge graph",
|
||||
name="search_nodes",
|
||||
description="Search for nodes",
|
||||
inputSchema=SearchNodesRequest.model_json_schema()
|
||||
),
|
||||
Tool(
|
||||
name="open_nodes",
|
||||
description="Open specific nodes by their names",
|
||||
description="Open specific nodes",
|
||||
inputSchema=OpenNodesRequest.model_json_schema()
|
||||
),
|
||||
Tool(
|
||||
name="add_observations",
|
||||
description="Add observations to existing entities",
|
||||
description="Add observations",
|
||||
inputSchema=AddObservationsRequest.model_json_schema()
|
||||
),
|
||||
Tool(
|
||||
name="create_relations",
|
||||
description="Create relations between entities",
|
||||
description="Create relations",
|
||||
inputSchema=CreateRelationsRequest.model_json_schema()
|
||||
),
|
||||
Tool(
|
||||
name="delete_entities",
|
||||
description="Delete entities from the knowledge graph",
|
||||
description="Delete entities",
|
||||
inputSchema=DeleteEntityRequest.model_json_schema()
|
||||
),
|
||||
Tool(
|
||||
name="delete_observations",
|
||||
description="Delete observations from entities",
|
||||
description="Delete observations",
|
||||
inputSchema=DeleteObservationsRequest.model_json_schema()
|
||||
)
|
||||
]
|
||||
logger.debug(f"Returning {len(tools)} available tools")
|
||||
return tools
|
||||
|
||||
|
||||
@self.call_tool()
|
||||
async def handle_call_tool(
|
||||
name: str,
|
||||
arguments: Dict[str, Any],
|
||||
*,
|
||||
memory_service: Optional[MemoryService] = None
|
||||
name: str,
|
||||
arguments: Dict[str, Any]
|
||||
) -> List[EmbeddedResource]:
|
||||
"""Handle tool calls by delegating to the memory service."""
|
||||
try:
|
||||
logger.debug(f"Handling tool call: {name} with args: {arguments}")
|
||||
# Check if tool exists
|
||||
if name not in TOOL_HANDLERS:
|
||||
logger.error(f"Unknown tool requested: {name}")
|
||||
raise McpError(METHOD_NOT_FOUND, f"Unknown tool: {name}")
|
||||
"""Forward tool calls to FastAPI endpoints."""
|
||||
|
||||
# Map tools to FastAPI endpoints
|
||||
handlers = {
|
||||
"create_entities": lambda c, a: c.post("/knowledge/entities", json=a),
|
||||
"search_nodes": lambda c, a: c.post("/knowledge/search", json=a),
|
||||
"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.delete(f"/knowledge/entities/{a['names'][0]}"),
|
||||
"delete_observations": lambda c, a: c.delete("/knowledge/observations", json=a)
|
||||
}
|
||||
|
||||
# Get tool endpoint
|
||||
handler = handlers.get(name)
|
||||
if handler is None:
|
||||
raise McpError(f"Unknown tool {name}")
|
||||
|
||||
# invoke the client handler function
|
||||
async with get_client() as client:
|
||||
response = await handler(client, arguments)
|
||||
return [create_response(response.json())]
|
||||
|
||||
|
||||
async with get_project_services(self.config.path) as service:
|
||||
tool_name = name # type: ignore
|
||||
result = [await TOOL_HANDLERS[tool_name](service, arguments)]
|
||||
logger.debug(f"Tool {name} completed successfully")
|
||||
return result
|
||||
|
||||
except ValueError as e:
|
||||
logger.error(f"Invalid parameters for {name}: {e}")
|
||||
raise McpError(INVALID_PARAMS, str(e))
|
||||
except EntityNotFoundError as e:
|
||||
logger.error(f"Entity not found in {name}: {e}")
|
||||
raise McpError(INVALID_PARAMS, str(e))
|
||||
except Exception as e:
|
||||
logger.exception(f"Unexpected error in {name}: {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
|
||||
logger.debug("Registered all handlers")
|
||||
|
||||
|
||||
# Create server instance with default config
|
||||
# Create server instance
|
||||
server = MemoryServer()
|
||||
|
||||
def setup_logging():
|
||||
"""Configure logging for the application."""
|
||||
# Remove default handler
|
||||
logger.remove()
|
||||
|
||||
# Add file handler
|
||||
logger.add(
|
||||
"basic-memory-mcp.log",
|
||||
rotation="100 MB",
|
||||
level="DEBUG",
|
||||
backtrace=True,
|
||||
diagnose=True
|
||||
)
|
||||
|
||||
# Add stdout handler for INFO and above
|
||||
logger.add(
|
||||
sys.stdout,
|
||||
level="INFO",
|
||||
backtrace=True,
|
||||
diagnose=True
|
||||
)
|
||||
|
||||
async def run_server():
|
||||
"""Run the MCP server."""
|
||||
from mcp.server.stdio import stdio_server
|
||||
|
||||
options = server.create_initialization_options()
|
||||
logger.info(f"Starting MCP server {options.server_name}")
|
||||
logger.info(f"Database URL: {server.config.database_url}")
|
||||
|
||||
async with stdio_server() as (read_stream, write_stream):
|
||||
await server.run(read_stream, write_stream, options)
|
||||
|
||||
if __name__ == "__main__":
|
||||
setup_logging()
|
||||
import asyncio
|
||||
from mcp.server.stdio import stdio_server
|
||||
|
||||
async def run_server():
|
||||
"""Run the MCP server."""
|
||||
options = server.create_initialization_options()
|
||||
async with stdio_server() as (read_stream, write_stream):
|
||||
await server.run(read_stream, write_stream, options)
|
||||
|
||||
asyncio.run(run_server())
|
||||
+131
-130
@@ -1,21 +1,43 @@
|
||||
"""Tests for the MCP server implementation."""
|
||||
"""Tests for the MCP server implementation using FastAPI TestClient."""
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from httpx import AsyncClient, ASGITransport
|
||||
|
||||
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.schemas import (
|
||||
CreateEntityResponse, SearchNodesResponse, AddObservationsResponse,
|
||||
)
|
||||
from basic_memory.api.app import app as fastapi_app
|
||||
from basic_memory.deps import get_project_config, get_engine
|
||||
from basic_memory.schemas import CreateEntityResponse, SearchNodesResponse, AddObservationsResponse
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def anyio_backend():
|
||||
return "asyncio"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app(test_config, engine) -> FastAPI:
|
||||
"""Create test FastAPI application."""
|
||||
app = fastapi_app
|
||||
app.dependency_overrides[get_project_config] = lambda: test_config
|
||||
app.dependency_overrides[get_engine] = lambda: engine
|
||||
return app
|
||||
|
||||
|
||||
@pytest.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.fixture
|
||||
def test_entity_data():
|
||||
"""Sample data for creating a test entity using camelCase (like MCP will)."""
|
||||
"""Sample data for creating a test entity."""
|
||||
return {
|
||||
"entities": [{
|
||||
"name": "Test Entity",
|
||||
@@ -25,14 +47,15 @@ def test_entity_data():
|
||||
}]
|
||||
}
|
||||
|
||||
|
||||
@pytest.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",
|
||||
"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",
|
||||
@@ -41,213 +64,191 @@ def test_directory_entity_data():
|
||||
}]
|
||||
}
|
||||
|
||||
@pytest.fixture
|
||||
def test_entity_snake_case():
|
||||
"""Same test data but using snake_case to test schema flexibility."""
|
||||
return {
|
||||
"entities": [{
|
||||
"name": "Test Entity",
|
||||
"entity_type": "test",
|
||||
"description": "", # Empty string instead of None
|
||||
"observations": ["This is a test observation"]
|
||||
}]
|
||||
}
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_list_tools(test_config):
|
||||
async def test_list_tools(app):
|
||||
"""Test that server exposes expected tools."""
|
||||
server_instance = MemoryServer(config=test_config)
|
||||
tools = await server_instance.handle_list_tools()
|
||||
|
||||
server = MemoryServer()
|
||||
tools = await server.handle_list_tools()
|
||||
|
||||
# Check each expected tool is present
|
||||
expected_tools = {
|
||||
"create_entities", "search_nodes", "open_nodes",
|
||||
"add_observations", "create_relations",
|
||||
"add_observations", "create_relations",
|
||||
"delete_entities", "delete_observations"
|
||||
}
|
||||
|
||||
|
||||
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.anyio
|
||||
async def test_create_directory_entity(test_directory_entity_data, memory_service, test_config):
|
||||
async def test_create_directory_entity(test_directory_entity_data, client):
|
||||
"""Test creating entity with exactly the data that failed in the tool."""
|
||||
server_instance = MemoryServer(config=test_config)
|
||||
result = await server_instance.handle_call_tool(
|
||||
"create_entities",
|
||||
test_directory_entity_data,
|
||||
memory_service=memory_service
|
||||
server = MemoryServer()
|
||||
result = await server.handle_call_tool(
|
||||
"create_entities",
|
||||
test_directory_entity_data
|
||||
)
|
||||
|
||||
|
||||
# Verify response format
|
||||
assert len(result) == 1
|
||||
assert isinstance(result[0], EmbeddedResource)
|
||||
assert result[0].type == "resource"
|
||||
|
||||
|
||||
# Verify entity creation
|
||||
response = CreateEntityResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue]
|
||||
response = CreateEntityResponse.model_validate_json(result[0].resource.text)
|
||||
assert len(response.entities) == 1
|
||||
assert response.entities[0].name == "Directory Organization"
|
||||
assert response.entities[0].entity_type == "memory"
|
||||
assert len(response.entities[0].observations) == 3
|
||||
created = response.entities[0]
|
||||
assert created.name == "Directory Organization"
|
||||
assert created.entity_type == "memory"
|
||||
assert len(created.observations) == 3
|
||||
|
||||
# Verify entity exists through API
|
||||
api_response = await client.get(f"/knowledge/entities/{created.id}")
|
||||
assert api_response.status_code == 200
|
||||
entity = api_response.json()
|
||||
assert entity["name"] == "Directory Organization"
|
||||
|
||||
# noinspection DuplicatedCode
|
||||
@pytest.mark.anyio
|
||||
async def test_create_entities_snake_case(test_entity_snake_case, memory_service, test_config):
|
||||
"""Test creating an entity with snake_case data (like internal usage)."""
|
||||
server_instance = MemoryServer(config=test_config)
|
||||
result = await server_instance.handle_call_tool(
|
||||
"create_entities",
|
||||
test_entity_snake_case,
|
||||
memory_service=memory_service
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
response = CreateEntityResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue]
|
||||
assert len(response.entities) == 1
|
||||
assert response.entities[0].name == "Test Entity"
|
||||
assert response.entities[0].entity_type == "test"
|
||||
assert len(response.entities[0].observations) == 1
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_search_nodes(test_entity_data, memory_service, test_config):
|
||||
async def test_search_nodes(test_entity_data, client):
|
||||
"""Test searching for an entity after creating it."""
|
||||
server_instance = MemoryServer(config=test_config)
|
||||
|
||||
server = MemoryServer()
|
||||
|
||||
# First create an entity
|
||||
await server_instance.handle_call_tool(
|
||||
"create_entities",
|
||||
test_entity_data,
|
||||
memory_service=memory_service
|
||||
)
|
||||
|
||||
await server.handle_call_tool("create_entities", test_entity_data)
|
||||
|
||||
# Then search for it
|
||||
result = await server_instance.handle_call_tool(
|
||||
"search_nodes",
|
||||
{"query": "Test Entity"},
|
||||
memory_service=memory_service
|
||||
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
|
||||
|
||||
response = SearchNodesResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue]
|
||||
|
||||
# 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.anyio
|
||||
async def test_add_observations(test_entity_data, memory_service, test_config):
|
||||
async def test_add_observations(test_entity_data, client):
|
||||
"""Test adding observations to an existing entity."""
|
||||
server_instance = MemoryServer(config=test_config)
|
||||
|
||||
# First create an entity and get its ID from response
|
||||
create_result = await server_instance.handle_call_tool(
|
||||
"create_entities",
|
||||
test_entity_data,
|
||||
memory_service=memory_service
|
||||
)
|
||||
|
||||
create_response = CreateEntityResponse.model_validate_json(create_result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue]
|
||||
server = MemoryServer()
|
||||
|
||||
# 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 observations using camelCase
|
||||
result = await server_instance.handle_call_tool(
|
||||
|
||||
# Add new observation
|
||||
result = await server.handle_call_tool(
|
||||
"add_observations",
|
||||
{
|
||||
"entity_id": entity_id,
|
||||
"observations": ["A new observation"]
|
||||
},
|
||||
memory_service=memory_service
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
# 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
|
||||
|
||||
response = AddObservationsResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue]
|
||||
|
||||
# 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.anyio
|
||||
async def test_invalid_tool_name(test_config):
|
||||
async def test_invalid_tool_name():
|
||||
"""Test calling a non-existent tool."""
|
||||
server_instance = MemoryServer(config=test_config)
|
||||
server = MemoryServer()
|
||||
with pytest.raises(McpError) as exc:
|
||||
await server_instance.handle_call_tool("not_a_tool", {})
|
||||
await server.handle_call_tool("not_a_tool", {})
|
||||
assert "Unknown tool" in str(exc.value)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
class TestInputValidation:
|
||||
"""Test input validation for various tools."""
|
||||
|
||||
async def test_missing_required_field(self, test_config):
|
||||
|
||||
async def test_missing_required_field(self):
|
||||
"""Test validation when required fields are missing."""
|
||||
server_instance = MemoryServer(config=test_config)
|
||||
|
||||
server = MemoryServer()
|
||||
|
||||
with pytest.raises(McpError) as exc:
|
||||
await server_instance.handle_call_tool("search_nodes", {})
|
||||
await server.handle_call_tool("search_nodes", {})
|
||||
assert "query" in str(exc.value).lower()
|
||||
|
||||
|
||||
with pytest.raises(McpError) as exc:
|
||||
await server_instance.handle_call_tool("create_entities", {})
|
||||
await server.handle_call_tool("create_entities", {})
|
||||
assert "entities" in str(exc.value).lower()
|
||||
|
||||
async def test_empty_arrays(self, test_config):
|
||||
|
||||
async def test_empty_arrays(self):
|
||||
"""Test validation of array fields that can't be empty."""
|
||||
server_instance = MemoryServer(config=test_config)
|
||||
|
||||
server = MemoryServer()
|
||||
|
||||
with pytest.raises(McpError) as exc:
|
||||
await server_instance.handle_call_tool("create_entities", {"entities": []})
|
||||
await server.handle_call_tool("create_entities", {"entities": []})
|
||||
assert "validation error" in str(exc.value).lower()
|
||||
|
||||
|
||||
with pytest.raises(McpError) as exc:
|
||||
await server_instance.handle_call_tool("open_nodes", {"names": []})
|
||||
await server.handle_call_tool("open_nodes", {"names": []})
|
||||
assert "validation error" in str(exc.value).lower()
|
||||
|
||||
async def test_invalid_field_types(self, test_config):
|
||||
|
||||
async def test_invalid_field_types(self):
|
||||
"""Test validation when fields have wrong types."""
|
||||
server_instance = MemoryServer(config=test_config)
|
||||
|
||||
server = MemoryServer()
|
||||
|
||||
with pytest.raises(McpError) as exc:
|
||||
await server_instance.handle_call_tool("search_nodes", {"query": 123})
|
||||
await server.handle_call_tool("search_nodes", {"query": 123})
|
||||
assert "str" in str(exc.value).lower()
|
||||
|
||||
|
||||
with pytest.raises(McpError) as exc:
|
||||
await server_instance.handle_call_tool("create_entities", {"entities": "not an array"})
|
||||
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()
|
||||
|
||||
async def test_invalid_nested_fields(self, test_config):
|
||||
|
||||
async def test_invalid_nested_fields(self):
|
||||
"""Test validation of nested object fields."""
|
||||
server_instance = MemoryServer(config=test_config)
|
||||
|
||||
server = MemoryServer()
|
||||
|
||||
with pytest.raises(McpError) as exc:
|
||||
await server_instance.handle_call_tool("create_entities", {
|
||||
await server.handle_call_tool("create_entities", {
|
||||
"entities": [{
|
||||
"name": "Test",
|
||||
# Missing required entity_type
|
||||
"observations": []
|
||||
}]
|
||||
})
|
||||
assert "entity_type" in str(exc.value).lower()
|
||||
|
||||
assert "entity_type" in str(exc.value).lower()
|
||||
Reference in New Issue
Block a user