mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
get_schema tool wip
This commit is contained in:
phernandez
@@ -11,8 +11,8 @@ from basic_memory.config import config
|
||||
from basic_memory.mcp.server import mcp
|
||||
|
||||
# Import tools to register them
|
||||
from basic_memory.mcp.tools import knowledge, search, documents
|
||||
__all__ = ["mcp", "knowledge", "search", "documents"]
|
||||
from basic_memory.mcp.tools import knowledge, search, documents, help
|
||||
__all__ = ["mcp", "knowledge", "search", "documents", "help"]
|
||||
|
||||
|
||||
def setup_logging(home_dir: str = config.home, log_file: str = "basic-memory.log"):
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
"""Enhanced MCP tool support with rich schema information."""
|
||||
|
||||
import inspect
|
||||
from typing import Any, Callable, Dict, List, Optional, Type
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class ToolExample(BaseModel):
|
||||
"""Example usage of a tool."""
|
||||
name: str
|
||||
description: str
|
||||
code: str
|
||||
|
||||
|
||||
class EnhancedToolMetadata(BaseModel):
|
||||
"""Enhanced tool metadata."""
|
||||
name: str = Field(description="Tool name")
|
||||
description: str = Field(description="Tool description")
|
||||
examples: List[ToolExample] = Field(
|
||||
default_factory=list,
|
||||
description="Example tool usage"
|
||||
)
|
||||
category: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Tool category for organization"
|
||||
)
|
||||
input_schema: Optional[Dict] = Field(
|
||||
default=None,
|
||||
description="Input parameter schema"
|
||||
)
|
||||
output_schema: Optional[Dict] = Field(
|
||||
default=None,
|
||||
description="Return value schema"
|
||||
)
|
||||
|
||||
|
||||
def enhanced_tool(
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
examples: Optional[List[Dict]] = None,
|
||||
category: Optional[str] = None,
|
||||
input_schema: Optional[Dict] = None,
|
||||
output_schema: Optional[Dict] = None,
|
||||
):
|
||||
"""Enhanced MCP tool decorator with rich metadata."""
|
||||
def decorator(fn: Callable):
|
||||
# Create metadata
|
||||
metadata = EnhancedToolMetadata(
|
||||
name=name or fn.__name__,
|
||||
description=description or fn.__doc__ or "",
|
||||
examples=[ToolExample(**ex) for ex in (examples or [])],
|
||||
category=category,
|
||||
input_schema=input_schema,
|
||||
output_schema=output_schema
|
||||
)
|
||||
|
||||
# Try to extract schemas from type hints if not provided
|
||||
if input_schema is None or output_schema is None:
|
||||
sig = inspect.signature(fn)
|
||||
|
||||
# Input schema from parameters
|
||||
if input_schema is None:
|
||||
for param in sig.parameters.values():
|
||||
if hasattr(param.annotation, "model_json_schema"):
|
||||
metadata.input_schema = param.annotation.model_json_schema()
|
||||
break
|
||||
|
||||
# Output schema from return type
|
||||
if output_schema is None:
|
||||
return_type = sig.return_annotation
|
||||
if hasattr(return_type, "model_json_schema"):
|
||||
metadata.output_schema = return_type.model_json_schema()
|
||||
|
||||
# Use regular MCP decorator first
|
||||
from basic_memory.mcp.server import mcp
|
||||
tool = mcp.tool(name=metadata.name, description=metadata.description)(fn)
|
||||
|
||||
# Store metadata on tool model
|
||||
setattr(tool, "_enhanced_metadata", metadata)
|
||||
|
||||
return tool
|
||||
|
||||
return decorator
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Help and schema tools for Basic Memory MCP server."""
|
||||
|
||||
from typing import Optional, Dict, List
|
||||
import inspect
|
||||
|
||||
from basic_memory.mcp.server import mcp
|
||||
from basic_memory.mcp.tools.enhanced import EnhancedToolMetadata
|
||||
|
||||
|
||||
def get_tool_metadata(tool) -> Optional[EnhancedToolMetadata]:
|
||||
"""Get enhanced tool metadata if available."""
|
||||
# Try to get metadata directly from tool
|
||||
if hasattr(tool, "_enhanced_metadata"):
|
||||
return getattr(tool, "_enhanced_metadata")
|
||||
return None
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def get_schema(tool_name: Optional[str] = None) -> Dict:
|
||||
"""Get schema information about available tools.
|
||||
|
||||
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
|
||||
- Example usage
|
||||
- Referenced Pydantic models
|
||||
|
||||
Examples:
|
||||
# Get complete tool catalog
|
||||
catalog = await get_schema()
|
||||
|
||||
# Get schema for specific tool
|
||||
entity_tool = await get_schema("create_entities")
|
||||
"""
|
||||
# Build catalog from enhanced tools
|
||||
catalog = {"tools": {}, "schemas": {}}
|
||||
|
||||
# Get all tools
|
||||
tools = await mcp.list_tools()
|
||||
|
||||
for tool in tools:
|
||||
# Get enhanced metadata if available
|
||||
metadata = get_tool_metadata(tool)
|
||||
if metadata:
|
||||
catalog["tools"][tool.name] = {
|
||||
"name": metadata.name,
|
||||
"description": metadata.description,
|
||||
"category": metadata.category,
|
||||
"examples": [ex.model_dump() for ex in metadata.examples],
|
||||
"inputSchema": metadata.input_schema,
|
||||
"outputSchema": metadata.output_schema,
|
||||
}
|
||||
else:
|
||||
# Basic metadata for non-enhanced tools
|
||||
model_data = tool.model_dump()
|
||||
catalog["tools"][tool.name] = {
|
||||
"name": tool.name,
|
||||
"description": tool.description,
|
||||
"inputSchema": model_data.get("inputSchema", {})
|
||||
}
|
||||
|
||||
if tool_name:
|
||||
if tool_name not in catalog["tools"]:
|
||||
raise ValueError(f"Unknown tool: {tool_name}")
|
||||
return {"tools": {tool_name: catalog["tools"][tool_name]}}
|
||||
|
||||
return catalog
|
||||
@@ -1,8 +1,6 @@
|
||||
"""Knowledge graph management tools for Basic Memory MCP server."""
|
||||
|
||||
from typing import Dict
|
||||
|
||||
import httpx
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from basic_memory.schemas.base import Entity, Relation, ObservationCategory, PathId
|
||||
from basic_memory.schemas.request import (
|
||||
@@ -17,47 +15,82 @@ from basic_memory.schemas.delete import (
|
||||
)
|
||||
from basic_memory.schemas.response import EntityListResponse, EntityResponse
|
||||
from basic_memory.mcp.async_client import client
|
||||
from basic_memory.mcp.server import mcp
|
||||
from basic_memory.services.exceptions import EntityNotFoundError
|
||||
from basic_memory.mcp.tools.enhanced import enhanced_tool
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def get_entity(path_id: PathId) -> EntityResponse:
|
||||
"""Get a specific entity by its path_id.
|
||||
@enhanced_tool(
|
||||
category="knowledge",
|
||||
examples=[{
|
||||
"name": "Create Component",
|
||||
"description": "Create a new technical component",
|
||||
"code": """
|
||||
await create_entities({
|
||||
"entities": [{
|
||||
"name": "SearchService",
|
||||
"entity_type": "component",
|
||||
"description": "Full-text search capability",
|
||||
"observations": [
|
||||
"Implements FTS5 for better performance",
|
||||
"Supports fuzzy matching"
|
||||
]
|
||||
}]
|
||||
})
|
||||
"""
|
||||
}]
|
||||
)
|
||||
async def create_entities(request: CreateEntityRequest) -> EntityListResponse:
|
||||
"""Create new entities in the knowledge graph.
|
||||
|
||||
Examples:
|
||||
# Load implementation details
|
||||
response = await get_entity("component/memory_service")
|
||||
|
||||
# Response contains complete entity:
|
||||
# EntityResponse(
|
||||
# path_id="component/memory_service",
|
||||
# name="memory_service",
|
||||
# entity_type="component",
|
||||
# description="Core knowledge persistence service",
|
||||
# observations=[
|
||||
# Observation(
|
||||
# category="TECH",
|
||||
# content="Using SQLite for storage",
|
||||
# context="Initial implementation"
|
||||
# ),
|
||||
# ...
|
||||
# ],
|
||||
# relations=[
|
||||
# Relation(
|
||||
# from_id="component/memory_service",
|
||||
# to_id="component/file_service",
|
||||
# relation_type="depends_on"
|
||||
# ),
|
||||
# ...
|
||||
# ]
|
||||
# )
|
||||
|
||||
# Load and analyze a design spec
|
||||
spec = await get_entity("specification/file_format")
|
||||
decisions = [obs for obs in spec.observations
|
||||
if obs.category == ObservationCategory.DESIGN]
|
||||
Entities can include initial observations and properties. Entity IDs
|
||||
are automatically generated from the type and name.
|
||||
"""
|
||||
url = "/knowledge/entities"
|
||||
response = await client.post(url, json=request.model_dump())
|
||||
return EntityListResponse.model_validate(response.json())
|
||||
|
||||
|
||||
@enhanced_tool(
|
||||
category="knowledge",
|
||||
examples=[{
|
||||
"name": "Add Dependency",
|
||||
"description": "Create dependency relationship between components",
|
||||
"code": """
|
||||
await create_relations({
|
||||
"relations": [{
|
||||
"from_id": "component/search_service",
|
||||
"to_id": "component/storage_service",
|
||||
"relation_type": "depends_on",
|
||||
"context": "Needs storage for search indexes"
|
||||
}]
|
||||
})
|
||||
"""
|
||||
}]
|
||||
)
|
||||
async def create_relations(request: CreateRelationsRequest) -> EntityListResponse:
|
||||
"""Create relations between existing entities."""
|
||||
url = "/knowledge/relations"
|
||||
response = await client.post(url, json=request.model_dump())
|
||||
return EntityListResponse.model_validate(response.json())
|
||||
|
||||
|
||||
@enhanced_tool(
|
||||
category="knowledge",
|
||||
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."""
|
||||
try:
|
||||
url = f"/knowledge/entities/{path_id}"
|
||||
response = await client.get(url)
|
||||
@@ -65,101 +98,14 @@ async def get_entity(path_id: PathId) -> EntityResponse:
|
||||
raise EntityNotFoundError(f"Entity not found: {path_id}")
|
||||
response.raise_for_status()
|
||||
return EntityResponse.model_validate(response.json())
|
||||
except httpx.HTTPStatusError as e:
|
||||
# If we got a 404, the entity doesn't exist
|
||||
if e.response.status_code == 404:
|
||||
except Exception as e:
|
||||
if hasattr(e, "response") and e.response.status_code == 404:
|
||||
raise EntityNotFoundError(f"Entity not found: {path_id}")
|
||||
# For any other HTTP error, re-raise
|
||||
raise
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def create_entities(request: CreateEntityRequest) -> EntityListResponse:
|
||||
"""Create new entities in the knowledge graph.
|
||||
|
||||
Examples:
|
||||
# Create a component with implementation details
|
||||
request = CreateEntityRequest(
|
||||
entities=[
|
||||
Entity(
|
||||
name="memory_service",
|
||||
entity_type="component",
|
||||
description="Core service for knowledge persistence",
|
||||
observations=[
|
||||
"Using SQLite for storage",
|
||||
"Implements filesystem as source of truth",
|
||||
"Handles atomic file operations"
|
||||
]
|
||||
)
|
||||
]
|
||||
)
|
||||
response = await create_entities(request)
|
||||
|
||||
# Response contains full entity details:
|
||||
# EntityListResponse(
|
||||
# entities=[
|
||||
# EntityResponse(
|
||||
# path_id="component/memory_service",
|
||||
# name="memory_service",
|
||||
# entity_type="component",
|
||||
# description="Core service for knowledge persistence",
|
||||
# observations=[...], # List[Observation]
|
||||
# relations=[] # Empty for new entities
|
||||
# )
|
||||
# ]
|
||||
# )
|
||||
"""
|
||||
url = "/knowledge/entities"
|
||||
response = await client.post(url, json=request.model_dump())
|
||||
return EntityListResponse.model_validate(response.json())
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def create_relations(request: CreateRelationsRequest) -> EntityListResponse:
|
||||
"""Create relations between existing entities.
|
||||
|
||||
Examples:
|
||||
# Document system dependencies
|
||||
request = CreateRelationsRequest(
|
||||
relations=[
|
||||
Relation(
|
||||
from_id="component/memory_service",
|
||||
to_id="component/file_service",
|
||||
relation_type="depends_on",
|
||||
context="File operations for persistence"
|
||||
),
|
||||
Relation(
|
||||
from_id="component/file_service",
|
||||
to_id="component/memory_service",
|
||||
relation_type="supports",
|
||||
context="Provides atomic file operations"
|
||||
)
|
||||
]
|
||||
)
|
||||
response = await create_relations(request)
|
||||
|
||||
# Response shows both entities with new relations:
|
||||
# EntityListResponse(
|
||||
# entities=[
|
||||
# EntityResponse( # memory_service
|
||||
# relations=[
|
||||
# Relation(to_id="component/file_service", ...)
|
||||
# ]
|
||||
# ),
|
||||
# EntityResponse( # file_service
|
||||
# relations=[
|
||||
# Relation(to_id="component/memory_service", ...)
|
||||
# ]
|
||||
# )
|
||||
# ]
|
||||
# )
|
||||
"""
|
||||
url = "/knowledge/relations"
|
||||
response = await client.post(url, json=request.model_dump())
|
||||
return EntityListResponse.model_validate(response.json())
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
@enhanced_tool()
|
||||
async def add_observations(request: AddObservationsRequest) -> EntityResponse:
|
||||
"""Add observations to an existing entity.
|
||||
|
||||
@@ -203,7 +149,7 @@ async def add_observations(request: AddObservationsRequest) -> EntityResponse:
|
||||
return EntityResponse.model_validate(response.json())
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
@enhanced_tool()
|
||||
async def delete_observations(request: DeleteObservationsRequest) -> EntityResponse:
|
||||
"""Delete specific observations from an entity.
|
||||
|
||||
@@ -229,7 +175,7 @@ async def delete_observations(request: DeleteObservationsRequest) -> EntityRespo
|
||||
return EntityResponse.model_validate(response.json())
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
@enhanced_tool()
|
||||
async def delete_relations(request: DeleteRelationsRequest) -> EntityListResponse:
|
||||
"""Delete relations between entities.
|
||||
|
||||
@@ -263,7 +209,7 @@ async def delete_relations(request: DeleteRelationsRequest) -> EntityListRespons
|
||||
return EntityListResponse.model_validate(response.json())
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
@enhanced_tool()
|
||||
async def delete_entities(request: DeleteEntitiesRequest) -> Dict[str, bool]:
|
||||
"""Delete entities from the knowledge graph.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user