fix tests

This commit is contained in:
phernandez
2024-12-30 23:42:57 -06:00
parent 31e81c9016
commit 705ca3e4c8
9 changed files with 255 additions and 259 deletions
+153 -8
View File
@@ -1,12 +1,157 @@
"""Shared MCP instance for Basic Memory."""
"""Enhanced FastMCP server instance for Basic Memory."""
from typing import Any, Callable, Dict, List, Optional
from fastmcp import FastMCP
from basic_memory.mcp.tools.enhanced import EnhancedToolManager
import inspect
from pydantic import BaseModel, Field
# Create and configure the shared MCP instance
mcp = FastMCP("Basic Memory")
from fastmcp.tools import Tool as FastMCPTool
from fastmcp.tools.tool_manager import ToolManager as FastMCPToolManager
# Replace the default tool manager with our enhanced version
mcp._tool_manager = EnhancedToolManager(
warn_on_duplicate_tools=mcp.settings.warn_on_duplicate_tools
)
class BasicMemoryServer(FastMCP):
"""Enhanced FastMCP server with schema support."""
def __init__(self, name: str | None = None, **settings: Any):
super().__init__(name=name, **settings)
# Replace default tool manager with our enhanced version
self._tool_manager = EnhancedToolManager(
warn_on_duplicate_tools=self.settings.warn_on_duplicate_tools
)
def tool(
self,
name: Optional[str] = None,
description: Optional[str] = None,
examples: Optional[List[Dict]] = None,
category: Optional[str] = None,
):
"""Decorator to register an enhanced tool.
Example:
@server.tool(
name="search",
description="Search for entities",
category="core",
examples=[{
"name": "Basic Search",
"description": "Search by text",
"code": 'results = await search({"query": "test"})'
}]
)
async def search(request: SearchRequest) -> SearchResults:
return await search_service.search(request)
"""
def decorator(fn: Callable) -> Callable:
tool = self._tool_manager.add_tool(
fn, name=name, description=description, examples=examples, category=category
)
return fn
return decorator
class ToolExample(BaseModel):
"""Example usage of a tool."""
name: str = Field(description="Name of the example")
description: str = Field(description="Description of what the example demonstrates")
code: str = Field(description="Example code")
class EnhancedTool(FastMCPTool):
"""Extended tool registration with rich metadata."""
examples: List[ToolExample] = Field(default_factory=list)
category: Optional[str] = Field(None)
input_schema: Optional[Dict] = Field(None)
output_schema: Optional[Dict] = Field(None)
@classmethod
def from_function(
cls,
fn: Callable,
name: Optional[str] = None,
description: Optional[str] = None,
examples: Optional[List[Dict]] = None,
category: Optional[str] = None,
) -> "EnhancedTool":
"""Create an enhanced tool from a function."""
# First create the base tool
base_tool = super().from_function(fn, name=name, description=description)
# Extract return type schema if available
return_schema = None
sig = inspect.signature(fn)
return_type = sig.return_annotation
if hasattr(return_type, "model_json_schema"):
return_schema = return_type.model_json_schema()
# Convert examples to ToolExample models
tool_examples = [ToolExample(**ex) for ex in (examples or [])]
return cls(
fn=fn,
name=base_tool.name,
description=base_tool.description,
parameters=base_tool.parameters,
fn_metadata=base_tool.fn_metadata,
is_async=base_tool.is_async,
context_kwarg=base_tool.context_kwarg,
examples=tool_examples,
category=category,
input_schema=base_tool.parameters,
output_schema=return_schema,
)
def get_schema(self) -> Dict:
"""Get complete tool schema including examples and metadata."""
return {
"name": self.name,
"description": self.description,
"category": self.category,
"inputSchema": self.input_schema,
"outputSchema": self.output_schema,
"examples": [ex.model_dump() for ex in self.examples],
}
class EnhancedToolManager(FastMCPToolManager):
"""Tool manager with enhanced metadata support."""
def add_tool(
self,
fn: Callable,
name: Optional[str] = None,
description: Optional[str] = None,
examples: Optional[List[Dict]] = None,
category: Optional[str] = None,
) -> EnhancedTool:
"""Add a tool with enhanced metadata."""
tool = EnhancedTool.from_function(
fn, name=name, description=description, examples=examples, category=category
)
self._tools[tool.name] = tool
return tool
def get_schema_catalog(self) -> Dict:
"""Get complete schema catalog for all tools."""
catalog = {"tools": {}, "categories": {}}
for tool in self._tools.values():
if isinstance(tool, EnhancedTool):
catalog["tools"][tool.name] = tool.get_schema()
if tool.category:
if tool.category not in catalog["categories"]:
catalog["categories"][tool.category] = {"name": tool.category, "tools": []}
catalog["categories"][tool.category]["tools"].append(tool.name)
return catalog
# Create the shared server instance
mcp = BasicMemoryServer("Basic Memory")
+2 -2
View File
@@ -5,11 +5,11 @@ from typing import List, Optional
from loguru import logger
from basic_memory.mcp.async_client import client
from basic_memory.mcp.tools.enhanced import enhanced_tool
from basic_memory.mcp.server import mcp
from basic_memory.schemas.activity import ActivityType, RecentActivity
@enhanced_tool()
@mcp.tool()
async def get_recent_activity(
timeframe: str = "1d",
activity_types: Optional[List[ActivityType]] = None,
+4 -4
View File
@@ -4,12 +4,12 @@ from typing import List, Optional
from loguru import logger
from basic_memory.mcp.tools.enhanced import enhanced_tool
from basic_memory.mcp.server import mcp
from basic_memory.schemas import EntityTypeList, ObservationCategoryList, TypedEntityList
from basic_memory.mcp.async_client import client
@enhanced_tool()
@mcp.tool()
async def get_entity_types() -> List[str]:
"""List all unique entity types in use across the knowledge graph.
@@ -33,7 +33,7 @@ async def get_entity_types() -> List[str]:
return EntityTypeList.model_validate(response.json())
@enhanced_tool()
@mcp.tool()
async def get_observation_categories() -> List[str]:
"""List all unique observation categories in use across the knowledge graph.
@@ -57,7 +57,7 @@ async def get_observation_categories() -> List[str]:
return ObservationCategoryList.model_validate(response.json())
@enhanced_tool()
@mcp.tool()
async def list_by_type(
entity_type: str, include_related: bool = False, sort_by: Optional[str] = "updated_at"
) -> TypedEntityList:
+21 -21
View File
@@ -2,25 +2,25 @@
from typing import Dict, List
from basic_memory.mcp.tools.enhanced import enhanced_tool
from basic_memory.mcp.server import mcp
from basic_memory.schemas.request import DocumentRequest, DocumentPathId
from basic_memory.schemas.response import DocumentResponse, DocumentCreateResponse
from basic_memory.mcp.async_client import client
@enhanced_tool()
@mcp.tool()
async def create_document(request: DocumentRequest) -> DocumentCreateResponse:
"""Create a new markdown document.
Examples:
# Create a technical specification
request = DocumentRequest(
path="specs/memory_format.md",
content='''# Memory Format Specification
## Overview
This document defines the standard format for memory files.
## Format
- Markdown with frontmatter
- UTF-8 encoding
@@ -33,7 +33,7 @@ async def create_document(request: DocumentRequest) -> DocumentCreateResponse:
}
)
response = await create_document(request)
# Response contains document info:
# DocumentCreateResponse(
# path="specs/memory_format.md",
@@ -48,16 +48,16 @@ async def create_document(request: DocumentRequest) -> DocumentCreateResponse:
return DocumentCreateResponse.model_validate(response.json())
@enhanced_tool()
@mcp.tool()
async def update_document(request: DocumentRequest) -> DocumentResponse:
"""Update an existing document.
Examples:
# Update implementation docs with new details
request = DocumentRequest(
path="docs/implementation.md",
content='''# Implementation Details
## Recent Changes
- Added FTS5 support
- Improved error handling
@@ -69,14 +69,14 @@ async def update_document(request: DocumentRequest) -> DocumentResponse:
}
)
response = await update_document(request)
# Response contains updated document:
# DocumentResponse(
# path="docs/implementation.md",
# content="# Implementation Details\n...",
# checksum="def456...",
# doc_metadata={...},
# created_at="2024-12-20T10:00:00Z",
# created_at="2024-12-20T10:00:00Z",
# updated_at="2024-12-25T14:30:00Z"
# )
"""
@@ -85,14 +85,14 @@ async def update_document(request: DocumentRequest) -> DocumentResponse:
return DocumentResponse.model_validate(response.json())
@enhanced_tool()
@mcp.tool()
async def get_document(path: DocumentPathId) -> DocumentResponse:
"""Get a document by its path.
Examples:
# Load an API specification
response = await get_document("specs/api_format.md")
# Response contains complete document:
# DocumentResponse(
# path="specs/api_format.md",
@@ -114,14 +114,14 @@ async def get_document(path: DocumentPathId) -> DocumentResponse:
return DocumentResponse.model_validate(response.json())
@enhanced_tool()
@mcp.tool()
async def list_documents() -> List[DocumentCreateResponse]:
"""List all documents in the system.
Examples:
# Get all documents with metadata
documents = await list_documents()
# Response is list of document info:
# [
# DocumentCreateResponse(
@@ -145,14 +145,14 @@ async def list_documents() -> List[DocumentCreateResponse]:
return [DocumentCreateResponse.model_validate(doc) for doc in response.json()]
@enhanced_tool()
@mcp.tool()
async def delete_document(path: DocumentPathId) -> Dict[str, bool]:
"""Delete a document.
Examples:
# Remove an obsolete document
result = await delete_document("docs/outdated_spec.md")
# Response indicates success:
# {
# "deleted": true
@@ -162,4 +162,4 @@ async def delete_document(path: DocumentPathId) -> Dict[str, bool]:
response = await client.delete(url)
if response.status_code == 204:
return {"deleted": True}
return response.json()
return response.json()
-137
View File
@@ -1,137 +0,0 @@
"""Enhanced MCP tool support with rich schema information."""
from typing import Any, Callable, Dict, List, Optional, Type
import inspect
from pydantic import BaseModel, Field
from fastmcp.tools import Tool as FastMCPTool
from fastmcp.tools.tool_manager import ToolManager as FastMCPToolManager
class ToolExample(BaseModel):
"""Example usage of a tool."""
name: str = Field(description="Name of the example")
description: str = Field(description="Description of what the example demonstrates")
code: str = Field(description="Example code")
class EnhancedTool(FastMCPTool):
"""Extended tool registration with rich metadata."""
examples: List[ToolExample] = Field(default_factory=list)
category: Optional[str] = Field(None)
input_schema: Optional[Dict] = Field(None)
output_schema: Optional[Dict] = Field(None)
@classmethod
def from_function(
cls,
fn: Callable,
name: Optional[str] = None,
description: Optional[str] = None,
examples: Optional[List[Dict]] = None,
category: Optional[str] = None,
) -> "EnhancedTool":
"""Create an enhanced tool from a function."""
# First create the base tool
base_tool = super().from_function(fn, name=name, description=description)
# Extract return type schema if available
return_schema = None
sig = inspect.signature(fn)
return_type = sig.return_annotation
if hasattr(return_type, "model_json_schema"):
return_schema = return_type.model_json_schema()
# Convert examples to ToolExample models
tool_examples = [ToolExample(**ex) for ex in (examples or [])]
return cls(
fn=fn,
name=base_tool.name,
description=base_tool.description,
parameters=base_tool.parameters,
fn_metadata=base_tool.fn_metadata,
is_async=base_tool.is_async,
context_kwarg=base_tool.context_kwarg,
examples=tool_examples,
category=category,
input_schema=base_tool.parameters,
output_schema=return_schema
)
def get_schema(self) -> Dict:
"""Get complete tool schema including examples and metadata."""
return {
"name": self.name,
"description": self.description,
"category": self.category,
"inputSchema": self.input_schema,
"outputSchema": self.output_schema,
"examples": [ex.model_dump() for ex in self.examples]
}
class EnhancedToolManager(FastMCPToolManager):
"""Tool manager with enhanced metadata support."""
def add_tool(
self,
fn: Callable,
name: Optional[str] = None,
description: Optional[str] = None,
examples: Optional[List[Dict]] = None,
category: Optional[str] = None,
) -> EnhancedTool:
"""Add a tool with enhanced metadata."""
tool = EnhancedTool.from_function(
fn,
name=name,
description=description,
examples=examples,
category=category
)
self._tools[tool.name] = tool
return tool
def get_schema_catalog(self) -> Dict:
"""Get complete schema catalog for all tools."""
catalog = {
"tools": {},
"categories": {}
}
for tool in self._tools.values():
if isinstance(tool, EnhancedTool):
catalog["tools"][tool.name] = tool.get_schema()
if tool.category:
if tool.category not in catalog["categories"]:
catalog["categories"][tool.category] = {
"name": tool.category,
"tools": []
}
catalog["categories"][tool.category]["tools"].append(tool.name)
return catalog
def enhanced_tool(
name: Optional[str] = None,
description: Optional[str] = None,
examples: Optional[List[Dict]] = None,
category: Optional[str] = None,
):
"""Decorator for registering enhanced tools."""
def decorator(fn: Callable) -> Callable:
# Store metadata on the function for later
if not hasattr(fn, "_tool_metadata"):
fn._tool_metadata = {}
fn._tool_metadata.update({
"name": name,
"description": description,
"examples": examples,
"category": category
})
return fn
return decorator
+18 -21
View File
@@ -3,40 +3,37 @@
from typing import Dict, Optional
from basic_memory.mcp.server import mcp
from basic_memory.mcp.tools.enhanced import enhanced_tool
@enhanced_tool(
@mcp.tool(
category="system",
examples=[
{
"name": "Get All Tools",
"description": "Get complete schema catalog for all tools",
"code": "catalog = await get_schema()"
"code": "catalog = await get_schema()",
},
{
"name": "Get Specific Tool",
"description": "Get schema for a specific tool",
"code": 'tool_schema = await get_schema("create_entity")'
}
]
"code": 'tool_schema = await get_schema("create_entity")',
},
],
)
async def get_schema(
tool_name: Optional[str] = None,
include_examples: bool = True,
include_referenced: bool = True
tool_name: Optional[str] = None, include_examples: bool = True, include_referenced: bool = True
) -> Dict:
"""Get schema information about available tools.
Args:
tool_name: Optional name of specific tool to get schema for
include_examples: Whether to include usage examples
include_referenced: Whether to include referenced model schemas
Returns:
Complete tool catalog if tool_name is None,
or specific tool schema if tool_name is provided.
Tool catalog includes:
- Tool descriptions and purposes
- Input/output schemas
@@ -46,30 +43,30 @@ async def get_schema(
"""
# Our tool manager has the enhanced schema support
catalog = mcp._tool_manager.get_schema_catalog()
# Filter if specific tool requested
if tool_name:
if tool_name not in catalog["tools"]:
raise ValueError(f"Unknown tool: {tool_name}")
tool_schema = catalog["tools"][tool_name]
if not include_examples:
tool_schema.pop("examples", None)
if include_referenced:
return {
"tools": {tool_name: tool_schema},
"referencedModels": tool_schema.get("referencedModels", {})
"referencedModels": tool_schema.get("referencedModels", {}),
}
else:
return {"tools": {tool_name: tool_schema}}
# Return full catalog with requested inclusions
result = catalog
if not include_examples:
for tool in result["tools"].values():
tool.pop("examples", None)
return result
return result
+47 -43
View File
@@ -4,7 +4,8 @@ from typing import Dict
import httpx
from basic_memory.schemas.base import Entity, Relation, ObservationCategory, PathId
from basic_memory.mcp.server import mcp
from basic_memory.schemas.base import PathId
from basic_memory.schemas.request import (
CreateEntityRequest,
CreateRelationsRequest,
@@ -13,20 +14,20 @@ from basic_memory.schemas.request import (
from basic_memory.schemas.delete import (
DeleteEntitiesRequest,
DeleteObservationsRequest,
DeleteRelationsRequest
DeleteRelationsRequest,
)
from basic_memory.schemas.response import EntityListResponse, EntityResponse
from basic_memory.mcp.async_client import client
from basic_memory.services.exceptions import EntityNotFoundError
from basic_memory.mcp.tools.enhanced import enhanced_tool
@enhanced_tool(
@mcp.tool(
category="knowledge",
examples=[{
"name": "Create Component",
"description": "Create a new technical component",
"code": """
examples=[
{
"name": "Create Component",
"description": "Create a new technical component",
"code": """
await create_entities({
"entities": [{
"name": "SearchService",
@@ -38,12 +39,13 @@ await create_entities({
]
}]
})
"""
}]
""",
}
],
)
async def create_entities(request: CreateEntityRequest) -> EntityListResponse:
"""Create new entities in the knowledge graph.
Entities can include initial observations and properties. Entity IDs
are automatically generated from the type and name.
"""
@@ -52,12 +54,13 @@ async def create_entities(request: CreateEntityRequest) -> EntityListResponse:
return EntityListResponse.model_validate(response.json())
@enhanced_tool(
@mcp.tool(
category="knowledge",
examples=[{
"name": "Add Dependency",
"description": "Create dependency relationship between components",
"code": """
examples=[
{
"name": "Add Dependency",
"description": "Create dependency relationship between components",
"code": """
await create_relations({
"relations": [{
"from_id": "component/search_service",
@@ -66,8 +69,9 @@ await create_relations({
"context": "Needs storage for search indexes"
}]
})
"""
}]
""",
}
],
)
async def create_relations(request: CreateRelationsRequest) -> EntityListResponse:
"""Create relations between existing entities."""
@@ -76,28 +80,30 @@ async def create_relations(request: CreateRelationsRequest) -> EntityListRespons
return EntityListResponse.model_validate(response.json())
@enhanced_tool(
@mcp.tool(
category="knowledge",
examples=[{
"name": "Get Entity Details",
"description": "Load complete entity information",
"code": """
examples=[
{
"name": "Get Entity Details",
"description": "Load complete entity information",
"code": """
# Get component details
entity = await get_entity("component/search_service")
print(f"Name: {entity.name}")
print(f"Type: {entity.entity_type}")
for obs in entity.observations:
print(f"- {obs.content}")
"""
}]
""",
}
],
)
async def get_entity(path_id: PathId) -> EntityResponse:
"""Get a specific entity by its path_id.
Examples:
# Load implementation details
response = await get_entity("component/memory_service")
# Response contains complete entity:
# EntityResponse(
# path_id="component/memory_service",
@@ -124,7 +130,7 @@ async def get_entity(path_id: PathId) -> EntityResponse:
# Load and analyze a design spec
spec = await get_entity("specification/file_format")
decisions = [obs for obs in spec.observations
decisions = [obs for obs in spec.observations
if obs.category == ObservationCategory.DESIGN]
"""
try:
@@ -142,12 +148,10 @@ async def get_entity(path_id: PathId) -> EntityResponse:
raise
@enhanced_tool()
@mcp.tool()
async def add_observations(request: AddObservationsRequest) -> EntityResponse:
"""Add observations to an existing entity.
Examples:
# Document implementation decisions with context
request = AddObservationsRequest(
@@ -169,7 +173,7 @@ async def add_observations(request: AddObservationsRequest) -> EntityResponse:
]
)
response = await add_observations(request)
# Response shows entity with new observations:
# EntityResponse(
# path_id="component/search_service",
@@ -188,10 +192,10 @@ async def add_observations(request: AddObservationsRequest) -> EntityResponse:
return EntityResponse.model_validate(response.json())
@enhanced_tool()
@mcp.tool()
async def delete_observations(request: DeleteObservationsRequest) -> EntityResponse:
"""Delete specific observations from an entity.
Examples:
# Remove obsolete implementation notes
request = DeleteObservationsRequest(
@@ -202,7 +206,7 @@ async def delete_observations(request: DeleteObservationsRequest) -> EntityRespo
]
)
response = await delete_observations(request)
# Response shows entity with observations removed:
# EntityResponse(
# path_id="component/indexer",
@@ -214,10 +218,10 @@ async def delete_observations(request: DeleteObservationsRequest) -> EntityRespo
return EntityResponse.model_validate(response.json())
@enhanced_tool()
@mcp.tool()
async def delete_relations(request: DeleteRelationsRequest) -> EntityListResponse:
"""Delete relations between entities.
Examples:
# Remove obsolete dependency
request = DeleteRelationsRequest(
@@ -230,7 +234,7 @@ async def delete_relations(request: DeleteRelationsRequest) -> EntityListRespons
]
)
response = await delete_relations(request)
# Response shows updated entities:
# EntityListResponse(
# entities=[
@@ -248,10 +252,10 @@ async def delete_relations(request: DeleteRelationsRequest) -> EntityListRespons
return EntityListResponse.model_validate(response.json())
@enhanced_tool()
@mcp.tool()
async def delete_entities(request: DeleteEntitiesRequest) -> Dict[str, bool]:
"""Delete entities from the knowledge graph.
Examples:
# Remove obsolete components
request = DeleteEntitiesRequest(
@@ -261,7 +265,7 @@ async def delete_entities(request: DeleteEntitiesRequest) -> Dict[str, bool]:
]
)
response = await delete_entities(request)
# Response indicates success:
# {
# "deleted": true
@@ -269,4 +273,4 @@ async def delete_entities(request: DeleteEntitiesRequest) -> Dict[str, bool]:
"""
url = "/knowledge/entities/delete"
response = await client.post(url, json=request.model_dump())
return response.json()
return response.json()
+3 -3
View File
@@ -2,13 +2,13 @@
from typing import Dict
from basic_memory.mcp.tools.enhanced import enhanced_tool
from basic_memory.mcp.server import mcp
from basic_memory.schemas.request import SearchNodesRequest, OpenNodesRequest
from basic_memory.schemas.response import SearchNodesResponse, EntityResponse
from basic_memory.mcp.async_client import client
@enhanced_tool()
@mcp.tool()
async def search_nodes(request: SearchNodesRequest) -> SearchNodesResponse:
"""Search for entities in the knowledge graph.
@@ -51,7 +51,7 @@ async def search_nodes(request: SearchNodesRequest) -> SearchNodesResponse:
return SearchNodesResponse.model_validate(response.json())
@enhanced_tool()
@mcp.tool()
async def open_nodes(request: OpenNodesRequest) -> Dict[str, EntityResponse]:
"""Load multiple entities by their path_ids.