From b818d7778b7ec8c51bad7cb7a3ca3cb2badd2f1d Mon Sep 17 00:00:00 2001 From: phernandez Date: Mon, 30 Dec 2024 12:41:45 -0600 Subject: [PATCH] get_schema tool wip --- Makefile | 2 +- pyproject.toml | 3 +- src/basic_memory/mcp/main.py | 4 +- src/basic_memory/mcp/tools/enhanced.py | 83 ++++++++++ src/basic_memory/mcp/tools/help.py | 70 ++++++++ src/basic_memory/mcp/tools/knowledge.py | 210 +++++++++--------------- tests/mcp/test_schema.py | 108 ++++++++++++ 7 files changed, 344 insertions(+), 136 deletions(-) create mode 100644 src/basic_memory/mcp/tools/enhanced.py create mode 100644 src/basic_memory/mcp/tools/help.py create mode 100644 tests/mcp/test_schema.py diff --git a/Makefile b/Makefile index 30f1d804..ae20cd09 100644 --- a/Makefile +++ b/Makefile @@ -30,4 +30,4 @@ clean: run-dev: - fastmcp dev src/basic_memory/mcp/server.py \ No newline at end of file + uv run fastmcp dev src/basic_memory/mcp/main.py \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 981b56bb..bf72a709 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,7 @@ dependencies = [ "basic-foundation", "markdown-it-py>=3.0.0", "python-frontmatter>=1.1.0", - "fastmcp>=0.4.1", + "fastmcp", # Removed version constraint since we're using local version "rich>=13.9.4", ] @@ -64,6 +64,7 @@ dev-dependencies = [ [tool.uv.sources] basic-foundation = { path = "../basic-foundation", editable = true } +fastmcp = { path = "../fastmcp", editable = true } [tool.pyright] include = [ diff --git a/src/basic_memory/mcp/main.py b/src/basic_memory/mcp/main.py index 9c725b7e..66883823 100644 --- a/src/basic_memory/mcp/main.py +++ b/src/basic_memory/mcp/main.py @@ -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"): diff --git a/src/basic_memory/mcp/tools/enhanced.py b/src/basic_memory/mcp/tools/enhanced.py new file mode 100644 index 00000000..f284ec5d --- /dev/null +++ b/src/basic_memory/mcp/tools/enhanced.py @@ -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 \ No newline at end of file diff --git a/src/basic_memory/mcp/tools/help.py b/src/basic_memory/mcp/tools/help.py new file mode 100644 index 00000000..5af8b53a --- /dev/null +++ b/src/basic_memory/mcp/tools/help.py @@ -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 \ No newline at end of file diff --git a/src/basic_memory/mcp/tools/knowledge.py b/src/basic_memory/mcp/tools/knowledge.py index 63cc0de4..19c38174 100644 --- a/src/basic_memory/mcp/tools/knowledge.py +++ b/src/basic_memory/mcp/tools/knowledge.py @@ -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. diff --git a/tests/mcp/test_schema.py b/tests/mcp/test_schema.py new file mode 100644 index 00000000..2f63debb --- /dev/null +++ b/tests/mcp/test_schema.py @@ -0,0 +1,108 @@ +"""Tests for MCP schema and tool discovery.""" + +import pytest +from pydantic import BaseModel +from typing import List, Optional + +from basic_memory.mcp.server import mcp +from basic_memory.mcp.tools.enhanced import enhanced_tool, EnhancedToolMetadata, ToolExample +from basic_memory.mcp.tools.help import get_schema + + +class TestInput(BaseModel): + """Test input model.""" + name: str + value: int + + +class TestOutput(BaseModel): + """Test output model.""" + result: str + values: List[int] + + +# Create the enhanced metadata explicitly +test_metadata = EnhancedToolMetadata( + name="test_tool", + description="A test tool with enhanced metadata", + category="test", + examples=[ + ToolExample( + name="Basic Usage", + description="Simple example", + code="await test_tool({\"name\": \"test\", \"value\": 42})" + ) + ] +) + +# Test tool with enhanced metadata +@enhanced_tool( + name="test_tool", + description="A test tool with enhanced metadata", + category="test", + examples=[{ + "name": "Basic Usage", + "description": "Simple example", + "code": "await test_tool({\"name\": \"test\", \"value\": 42})" + }] +) +async def test_tool(request: TestInput) -> TestOutput: + """Test tool.""" + # Add metadata directly + setattr(test_tool, "_enhanced_metadata", test_metadata) + return TestOutput( + result=f"Processed {request.name}", + values=[request.value] + ) + + +# Test basic tool without enhancements +@mcp.tool(name="basic_tool") +async def basic_tool(value: str) -> str: + """A basic tool without enhanced metadata""" + return f"Echo: {value}" + + +@pytest.mark.asyncio +async def test_get_schema_all(): + """Test getting complete tool catalog.""" + catalog = await get_schema() + + assert "tools" in catalog + assert "test_tool" in catalog["tools"] + assert "basic_tool" in catalog["tools"] + + +@pytest.mark.asyncio +async def test_get_schema_enhanced_tool(): + """Test getting schema for enhanced tool.""" + schema = await get_schema("test_tool") + + assert "tools" in schema + assert "test_tool" in schema["tools"] + + tool_info = schema["tools"]["test_tool"] + assert tool_info["category"] == "test" + assert len(tool_info["examples"]) == 1 + assert tool_info["inputSchema"] is not None + + +@pytest.mark.asyncio +async def test_get_schema_basic_tool(): + """Test getting schema for basic tool.""" + schema = await get_schema("basic_tool") + + assert "tools" in schema + assert "basic_tool" in schema["tools"] + + tool_info = schema["tools"]["basic_tool"] + assert tool_info["name"] == "basic_tool" + assert tool_info["description"] == "A basic tool without enhanced metadata" + assert "inputSchema" in tool_info + + +@pytest.mark.asyncio +async def test_get_schema_unknown_tool(): + """Test getting schema for unknown tool.""" + with pytest.raises(ValueError, match="Unknown tool: unknown_tool"): + await get_schema("unknown_tool") \ No newline at end of file