494 lines
19 KiB
Python
494 lines
19 KiB
Python
"""Shared MCP tool implementations for Hindsight.
|
|
|
|
This module provides the core tool logic used by both:
|
|
- mcp_local.py (stdio transport for Claude Code)
|
|
- api/mcp.py (HTTP transport for API server)
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from typing import Any, Callable
|
|
|
|
from fastmcp import FastMCP
|
|
|
|
from hindsight_api import MemoryEngine
|
|
from hindsight_api.config import (
|
|
DEFAULT_MCP_RECALL_DESCRIPTION,
|
|
DEFAULT_MCP_RETAIN_DESCRIPTION,
|
|
)
|
|
from hindsight_api.engine.memory_engine import Budget
|
|
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
|
|
from hindsight_api.models import RequestContext
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class MCPToolsConfig:
|
|
"""Configuration for MCP tools registration."""
|
|
|
|
# How to resolve bank_id for operations
|
|
bank_id_resolver: Callable[[], str | None]
|
|
|
|
# Whether to include bank_id as a parameter on tools (for multi-bank support)
|
|
include_bank_id_param: bool = False
|
|
|
|
# Which tools to register
|
|
tools: set[str] | None = None # None means all tools
|
|
|
|
# Custom descriptions (if None, uses defaults)
|
|
retain_description: str | None = None
|
|
recall_description: str | None = None
|
|
|
|
# Retain behavior
|
|
retain_fire_and_forget: bool = False # If True, use asyncio.create_task pattern
|
|
|
|
|
|
def parse_timestamp(timestamp: str) -> datetime | None:
|
|
"""Parse an ISO format timestamp string.
|
|
|
|
Args:
|
|
timestamp: ISO format timestamp (e.g., '2024-01-15T10:30:00Z')
|
|
|
|
Returns:
|
|
Parsed datetime or None if invalid
|
|
|
|
Raises:
|
|
ValueError: If timestamp format is invalid
|
|
"""
|
|
try:
|
|
return datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
|
|
except ValueError as e:
|
|
raise ValueError(
|
|
f"Invalid timestamp format '{timestamp}'. "
|
|
"Expected ISO format like '2024-01-15T10:30:00' or '2024-01-15T10:30:00Z'"
|
|
) from e
|
|
|
|
|
|
def build_content_dict(
|
|
content: str,
|
|
context: str,
|
|
timestamp: str | None = None,
|
|
) -> tuple[dict[str, Any], str | None]:
|
|
"""Build a content dict for retain operations.
|
|
|
|
Args:
|
|
content: The memory content
|
|
context: Category for the memory
|
|
timestamp: Optional ISO timestamp
|
|
|
|
Returns:
|
|
Tuple of (content_dict, error_message). error_message is None if successful.
|
|
"""
|
|
content_dict: dict[str, Any] = {"content": content, "context": context}
|
|
|
|
if timestamp:
|
|
try:
|
|
parsed_timestamp = parse_timestamp(timestamp)
|
|
content_dict["event_date"] = parsed_timestamp
|
|
except ValueError as e:
|
|
return {}, str(e)
|
|
|
|
return content_dict, None
|
|
|
|
|
|
def register_mcp_tools(
|
|
mcp: FastMCP,
|
|
memory: MemoryEngine,
|
|
config: MCPToolsConfig,
|
|
) -> None:
|
|
"""Register MCP tools on a FastMCP server.
|
|
|
|
Args:
|
|
mcp: FastMCP server instance
|
|
memory: MemoryEngine instance
|
|
config: Tool configuration
|
|
"""
|
|
tools_to_register = config.tools or {"retain", "recall", "reflect", "list_banks", "create_bank"}
|
|
|
|
if "retain" in tools_to_register:
|
|
_register_retain(mcp, memory, config)
|
|
|
|
if "recall" in tools_to_register:
|
|
_register_recall(mcp, memory, config)
|
|
|
|
if "reflect" in tools_to_register:
|
|
_register_reflect(mcp, memory, config)
|
|
|
|
if "list_banks" in tools_to_register:
|
|
_register_list_banks(mcp, memory, config)
|
|
|
|
if "create_bank" in tools_to_register:
|
|
_register_create_bank(mcp, memory, config)
|
|
|
|
|
|
def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
|
"""Register the retain tool."""
|
|
description = config.retain_description or DEFAULT_MCP_RETAIN_DESCRIPTION
|
|
|
|
if config.include_bank_id_param:
|
|
if config.retain_fire_and_forget:
|
|
|
|
@mcp.tool(description=description)
|
|
async def retain(
|
|
content: str,
|
|
context: str = "general",
|
|
timestamp: str | None = None,
|
|
bank_id: str | None = None,
|
|
) -> dict:
|
|
"""
|
|
Args:
|
|
content: The fact/memory to store (be specific and include relevant details)
|
|
context: Category for the memory (e.g., 'preferences', 'work', 'hobbies', 'family'). Default: 'general'
|
|
timestamp: When this event/fact occurred (ISO format, e.g., '2024-01-15T10:30:00Z'). Useful for timeline tracking.
|
|
bank_id: Optional bank to store in (defaults to session bank). Use for cross-bank operations.
|
|
"""
|
|
import asyncio
|
|
|
|
target_bank = bank_id or config.bank_id_resolver()
|
|
if target_bank is None:
|
|
return {"status": "error", "message": "No bank_id configured"}
|
|
|
|
content_dict, error = build_content_dict(content, context, timestamp)
|
|
if error:
|
|
return {"status": "error", "message": error}
|
|
|
|
async def _retain():
|
|
try:
|
|
await memory.retain_batch_async(
|
|
bank_id=target_bank,
|
|
contents=[content_dict],
|
|
request_context=RequestContext(),
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Error storing memory: {e}", exc_info=True)
|
|
|
|
asyncio.create_task(_retain())
|
|
return {"status": "accepted", "message": "Memory storage initiated"}
|
|
|
|
else:
|
|
|
|
@mcp.tool(description=description)
|
|
async def retain(
|
|
content: str,
|
|
context: str = "general",
|
|
timestamp: str | None = None,
|
|
async_processing: bool = True,
|
|
bank_id: str | None = None,
|
|
) -> str:
|
|
"""
|
|
Args:
|
|
content: The fact/memory to store (be specific and include relevant details)
|
|
context: Category for the memory (e.g., 'preferences', 'work', 'hobbies', 'family'). Default: 'general'
|
|
timestamp: When this event/fact occurred (ISO format, e.g., '2024-01-15T10:30:00Z'). Useful for timeline tracking.
|
|
async_processing: If True, queue for background processing and return immediately. If False, wait for completion. Default: True
|
|
bank_id: Optional bank to store in (defaults to session bank). Use for cross-bank operations.
|
|
"""
|
|
try:
|
|
target_bank = bank_id or config.bank_id_resolver()
|
|
if target_bank is None:
|
|
return "Error: No bank_id configured"
|
|
|
|
content_dict, error = build_content_dict(content, context, timestamp)
|
|
if error:
|
|
return f"Error: {error}"
|
|
|
|
contents = [content_dict]
|
|
if async_processing:
|
|
result = await memory.submit_async_retain(
|
|
bank_id=target_bank, contents=contents, request_context=RequestContext()
|
|
)
|
|
return f"Memory queued for background processing (operation_id: {result.get('operation_id', 'N/A')})"
|
|
else:
|
|
await memory.retain_batch_async(
|
|
bank_id=target_bank,
|
|
contents=contents,
|
|
request_context=RequestContext(),
|
|
)
|
|
return f"Memory stored successfully in bank '{target_bank}'"
|
|
except Exception as e:
|
|
logger.error(f"Error storing memory: {e}", exc_info=True)
|
|
return f"Error: {str(e)}"
|
|
|
|
else:
|
|
# No bank_id param - use fixed bank from resolver
|
|
|
|
@mcp.tool(description=description)
|
|
async def retain(
|
|
content: str,
|
|
context: str = "general",
|
|
timestamp: str | None = None,
|
|
) -> dict:
|
|
"""
|
|
Args:
|
|
content: The fact/memory to store (be specific and include relevant details)
|
|
context: Category for the memory (e.g., 'preferences', 'work', 'hobbies', 'family'). Default: 'general'
|
|
timestamp: When this event/fact occurred (ISO format, e.g., '2024-01-15T10:30:00Z'). Useful for timeline tracking.
|
|
"""
|
|
import asyncio
|
|
|
|
target_bank = config.bank_id_resolver()
|
|
if target_bank is None:
|
|
return {"status": "error", "message": "No bank_id configured"}
|
|
|
|
content_dict, error = build_content_dict(content, context, timestamp)
|
|
if error:
|
|
return {"status": "error", "message": error}
|
|
|
|
async def _retain():
|
|
try:
|
|
await memory.retain_batch_async(
|
|
bank_id=target_bank,
|
|
contents=[content_dict],
|
|
request_context=RequestContext(),
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Error storing memory: {e}", exc_info=True)
|
|
|
|
asyncio.create_task(_retain())
|
|
return {"status": "accepted", "message": "Memory storage initiated"}
|
|
|
|
|
|
def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
|
"""Register the recall tool."""
|
|
description = config.recall_description or DEFAULT_MCP_RECALL_DESCRIPTION
|
|
|
|
if config.include_bank_id_param:
|
|
|
|
@mcp.tool(description=description)
|
|
async def recall(
|
|
query: str,
|
|
max_tokens: int = 4096,
|
|
bank_id: str | None = None,
|
|
) -> str | dict:
|
|
"""
|
|
Args:
|
|
query: Natural language search query (e.g., "user's food preferences", "what projects is user working on")
|
|
max_tokens: Maximum tokens to return in results (default: 4096)
|
|
bank_id: Optional bank to search in (defaults to session bank). Use for cross-bank operations.
|
|
"""
|
|
try:
|
|
target_bank = bank_id or config.bank_id_resolver()
|
|
if target_bank is None:
|
|
return "Error: No bank_id configured"
|
|
|
|
recall_result = await memory.recall_async(
|
|
bank_id=target_bank,
|
|
query=query,
|
|
fact_type=list(VALID_RECALL_FACT_TYPES),
|
|
budget=Budget.HIGH,
|
|
max_tokens=max_tokens,
|
|
request_context=RequestContext(),
|
|
)
|
|
|
|
return recall_result.model_dump_json(indent=2)
|
|
except Exception as e:
|
|
logger.error(f"Error searching: {e}", exc_info=True)
|
|
return f'{{"error": "{e}", "results": []}}'
|
|
|
|
else:
|
|
|
|
@mcp.tool(description=description)
|
|
async def recall(
|
|
query: str,
|
|
max_tokens: int = 4096,
|
|
) -> dict:
|
|
"""
|
|
Args:
|
|
query: Natural language search query (e.g., "user's food preferences", "what projects is user working on")
|
|
max_tokens: Maximum tokens to return in results (default: 4096)
|
|
"""
|
|
try:
|
|
target_bank = config.bank_id_resolver()
|
|
if target_bank is None:
|
|
return {"error": "No bank_id configured", "results": []}
|
|
|
|
recall_result = await memory.recall_async(
|
|
bank_id=target_bank,
|
|
query=query,
|
|
fact_type=list(VALID_RECALL_FACT_TYPES),
|
|
budget=Budget.HIGH,
|
|
max_tokens=max_tokens,
|
|
request_context=RequestContext(),
|
|
)
|
|
|
|
return recall_result.model_dump()
|
|
except Exception as e:
|
|
logger.error(f"Error searching: {e}", exc_info=True)
|
|
return {"error": str(e), "results": []}
|
|
|
|
|
|
def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
|
"""Register the reflect tool."""
|
|
|
|
if config.include_bank_id_param:
|
|
|
|
@mcp.tool()
|
|
async def reflect(
|
|
query: str,
|
|
context: str | None = None,
|
|
budget: str = "low",
|
|
bank_id: str | None = None,
|
|
) -> str:
|
|
"""
|
|
Generate thoughtful analysis by synthesizing stored memories with the bank's personality.
|
|
|
|
WHEN TO USE THIS TOOL:
|
|
Use reflect when you need reasoned analysis, not just fact retrieval. This tool
|
|
thinks through the question using everything the bank knows and its personality traits.
|
|
|
|
EXAMPLES OF GOOD QUERIES:
|
|
- "What patterns have emerged in how I approach debugging?"
|
|
- "Based on my past decisions, what architectural style do I prefer?"
|
|
- "What might be the best approach for this problem given what you know about me?"
|
|
- "How should I prioritize these tasks based on my goals?"
|
|
|
|
HOW IT DIFFERS FROM RECALL:
|
|
- recall: Returns raw facts matching your search (fast lookup)
|
|
- reflect: Reasons across memories to form a synthesized answer (deeper analysis)
|
|
|
|
Use recall for "what did I say about X?" and reflect for "what should I do about X?"
|
|
|
|
Args:
|
|
query: The question or topic to reflect on
|
|
context: Optional context about why this reflection is needed
|
|
budget: Search budget - 'low', 'mid', or 'high' (default: 'low')
|
|
bank_id: Optional bank to reflect in (defaults to session bank). Use for cross-bank operations.
|
|
"""
|
|
try:
|
|
target_bank = bank_id or config.bank_id_resolver()
|
|
if target_bank is None:
|
|
return "Error: No bank_id configured"
|
|
|
|
budget_map = {"low": Budget.LOW, "mid": Budget.MID, "high": Budget.HIGH}
|
|
budget_enum = budget_map.get(budget.lower(), Budget.LOW)
|
|
|
|
reflect_result = await memory.reflect_async(
|
|
bank_id=target_bank,
|
|
query=query,
|
|
budget=budget_enum,
|
|
context=context,
|
|
request_context=RequestContext(),
|
|
)
|
|
|
|
return reflect_result.model_dump_json(indent=2)
|
|
except Exception as e:
|
|
logger.error(f"Error reflecting: {e}", exc_info=True)
|
|
return f'{{"error": "{e}", "text": ""}}'
|
|
|
|
else:
|
|
|
|
@mcp.tool()
|
|
async def reflect(
|
|
query: str,
|
|
context: str | None = None,
|
|
budget: str = "low",
|
|
) -> dict:
|
|
"""
|
|
Generate thoughtful analysis by synthesizing stored memories with the bank's personality.
|
|
|
|
WHEN TO USE THIS TOOL:
|
|
Use reflect when you need reasoned analysis, not just fact retrieval. This tool
|
|
thinks through the question using everything the bank knows and its personality traits.
|
|
|
|
EXAMPLES OF GOOD QUERIES:
|
|
- "What patterns have emerged in how I approach debugging?"
|
|
- "Based on my past decisions, what architectural style do I prefer?"
|
|
- "What might be the best approach for this problem given what you know about me?"
|
|
- "How should I prioritize these tasks based on my goals?"
|
|
|
|
HOW IT DIFFERS FROM RECALL:
|
|
- recall: Returns raw facts matching your search (fast lookup)
|
|
- reflect: Reasons across memories to form a synthesized answer (deeper analysis)
|
|
|
|
Use recall for "what did I say about X?" and reflect for "what should I do about X?"
|
|
|
|
Args:
|
|
query: The question or topic to reflect on
|
|
context: Optional context about why this reflection is needed
|
|
budget: Search budget - 'low', 'mid', or 'high' (default: 'low')
|
|
"""
|
|
try:
|
|
target_bank = config.bank_id_resolver()
|
|
if target_bank is None:
|
|
return {"error": "No bank_id configured", "text": ""}
|
|
|
|
budget_map = {"low": Budget.LOW, "mid": Budget.MID, "high": Budget.HIGH}
|
|
budget_enum = budget_map.get(budget.lower(), Budget.LOW)
|
|
|
|
reflect_result = await memory.reflect_async(
|
|
bank_id=target_bank,
|
|
query=query,
|
|
budget=budget_enum,
|
|
context=context,
|
|
request_context=RequestContext(),
|
|
)
|
|
|
|
return reflect_result.model_dump()
|
|
except Exception as e:
|
|
logger.error(f"Error reflecting: {e}", exc_info=True)
|
|
return {"error": str(e), "text": ""}
|
|
|
|
|
|
def _register_list_banks(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
|
"""Register the list_banks tool."""
|
|
|
|
@mcp.tool()
|
|
async def list_banks() -> str:
|
|
"""
|
|
List all available memory banks.
|
|
|
|
Use this tool to discover what memory banks exist in the system.
|
|
Each bank is an isolated memory store (like a separate "brain").
|
|
|
|
Returns:
|
|
JSON list of banks with their IDs, names, dispositions, and missions.
|
|
"""
|
|
try:
|
|
banks = await memory.list_banks(request_context=RequestContext())
|
|
return json.dumps({"banks": banks}, indent=2)
|
|
except Exception as e:
|
|
logger.error(f"Error listing banks: {e}", exc_info=True)
|
|
return f'{{"error": "{e}", "banks": []}}'
|
|
|
|
|
|
def _register_create_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
|
"""Register the create_bank tool."""
|
|
|
|
@mcp.tool()
|
|
async def create_bank(bank_id: str, name: str | None = None, mission: str | None = None) -> str:
|
|
"""
|
|
Create a new memory bank or get an existing one.
|
|
|
|
Memory banks are isolated stores - each one is like a separate "brain" for a user/agent.
|
|
Banks are auto-created with default settings if they don't exist.
|
|
|
|
Args:
|
|
bank_id: Unique identifier for the bank (e.g., 'user-123', 'agent-alpha')
|
|
name: Optional human-friendly name for the bank
|
|
mission: Optional mission describing who the agent is and what they're trying to accomplish
|
|
"""
|
|
try:
|
|
# get_bank_profile auto-creates bank if it doesn't exist
|
|
profile = await memory.get_bank_profile(bank_id, request_context=RequestContext())
|
|
|
|
# Update name/mission if provided
|
|
if name is not None or mission is not None:
|
|
await memory.update_bank(
|
|
bank_id,
|
|
name=name,
|
|
mission=mission,
|
|
request_context=RequestContext(),
|
|
)
|
|
# Fetch updated profile
|
|
profile = await memory.get_bank_profile(bank_id, request_context=RequestContext())
|
|
|
|
# Serialize disposition if it's a Pydantic model
|
|
if "disposition" in profile and hasattr(profile["disposition"], "model_dump"):
|
|
profile["disposition"] = profile["disposition"].model_dump()
|
|
return json.dumps(profile, indent=2)
|
|
except Exception as e:
|
|
logger.error(f"Error creating bank: {e}", exc_info=True)
|
|
return f'{{"error": "{e}"}}'
|