"""Tools for managing memory optimization."""
from pathlib import Path
from typing import Annotated, Optional
from fastmcp import Context
from ..memory_optimizer import MemoryOptimizer
from ..server_registry import get_server_registry
def register_memory_tools() -> None:
"""Register all memory optimization tools with the server."""
registry = get_server_registry()
app = registry.app
instruction_manager = registry.instruction_manager
read_only = registry.read_only
@app.tool(
name="optimize_memory",
description="Manually optimize a memory file using AI to reorganize and consolidate entries while preserving all information.",
tags={"public", "memory"},
annotations={
"idempotentHint": False,
"readOnlyHint": False,
"title": "Optimize Memory File",
"parameters": {
"memory_file": "Optional path to specific memory file. If not provided, will optimize the user's main memory file.",
"force": "Force optimization even if criteria are not met. Defaults to False.",
},
"returns": "Returns detailed results of the optimization process including status, entries before/after, and backup information.",
},
meta={
"category": "memory",
},
)
async def optimize_memory(
ctx: Context,
memory_file: Annotated[Optional[str], "Path to memory file to optimize"] = None,
force: Annotated[bool, "Force optimization regardless of criteria"] = False,
) -> str:
"""Manually optimize a memory file using AI sampling."""
if read_only:
return "Error: Server is running in read-only mode"
try:
# Determine which file to optimize
if memory_file:
file_path = Path(memory_file)
if not file_path.exists():
return f"Error: Memory file not found: {memory_file}"
else:
# Use default user memory file
user_memory_path = instruction_manager.get_memory_file_path()
if not user_memory_path.exists():
return "Error: No user memory file found to optimize"
file_path = user_memory_path
# Create optimizer and run optimization
optimizer = MemoryOptimizer(instruction_manager)
result = await optimizer.optimize_memory_if_needed(file_path, ctx, force=force)
# Format result message
status = result.get("status", "unknown")
if status == "optimized":
entries_before = result.get("entries_before", "unknown")
entries_after = result.get("entries_after", "unknown")
backup_created = result.get("backup_created", False)
message = "✅ Memory optimization completed successfully!\n"
message += f"📊 Entries: {entries_before} → {entries_after}\n"
message += f"🔄 Method: {result.get('method', 'ai')}\n"
message += f"💾 Backup created: {'Yes' if backup_created else 'No'}\n"
message += f"📝 Reason: {result.get('reason', 'Manual optimization')}"
elif status == "metadata_updated":
message = "📝 Memory metadata updated (AI optimization unavailable)\n"
message += f"💾 Backup created: {'Yes' if result.get('backup_created', False) else 'No'}\n"
message += f"📝 Reason: {result.get('reason', 'Manual optimization')}"
elif status == "skipped":
message = f"⏭️ Optimization skipped: {result.get('reason', 'Unknown reason')}\n"
message += "💡 Use force=True to optimize anyway"
elif status == "error":
message = f"❌ Optimization failed: {result.get('reason', 'Unknown error')}"
else:
message = f"🔍 Optimization result: {status}"
return message
except Exception as e:
return f"Error during memory optimization: {str(e)}"
@app.tool(
name="memory_stats",
description="Get detailed statistics and optimization status for a memory file.",
tags={"public", "memory"},
annotations={
"idempotentHint": True,
"readOnlyHint": True,
"title": "Memory File Statistics",
"parameters": {
"memory_file": "Optional path to specific memory file. If not provided, will show stats for the user's main memory file.",
},
"returns": "Returns comprehensive statistics including file size, entry count, optimization eligibility, and configuration settings.",
},
meta={
"category": "memory",
},
)
def memory_stats(
memory_file: Annotated[Optional[str], "Path to memory file to analyze"] = None,
) -> str:
"""Get detailed statistics about a memory file."""
try:
# Determine which file to analyze
if memory_file:
file_path = Path(memory_file)
if not file_path.exists():
return f"Error: Memory file not found: {memory_file}"
else:
# Use default user memory file
user_memory_path = instruction_manager.get_memory_file_path()
if not user_memory_path.exists():
return "No user memory file found"
file_path = user_memory_path
# Get stats
optimizer = MemoryOptimizer(instruction_manager)
stats = optimizer.get_memory_stats(file_path)
if "error" in stats:
return str(stats["error"])
# Format stats message
message = "📊 **Memory File Statistics**\n\n"
message += f"📁 **File**: `{stats['file_path']}`\n"
message += f"📏 **Size**: {stats['file_size_bytes']:,} bytes\n"
message += f"🎯 **Tokens**: {stats['current_tokens']:,} (Last optimized: {stats['last_optimized_tokens']:,})\n"
message += f"📈 **Growth**: {stats['token_growth_percent']}% since last optimization\n"
message += f"📝 **Entries**: {stats['current_entries']}\n"
message += f"🔄 **Last Optimized**: {stats['last_optimized'] or 'Never'}\n"
message += f"⚡ **Optimization Version**: {stats['optimization_version']}\n\n"
message += "⚙️ **Configuration**:\n"
message += f"• Auto-optimize: {'✅ Enabled' if stats['auto_optimize_enabled'] else '❌ Disabled'}\n"
threshold_percent = int((stats["token_growth_threshold"] - 1.0) * 100)
message += f"• Token growth threshold: {threshold_percent}% ({stats['token_growth_threshold']})\n\n"
message += "🎯 **Optimization Status**:\n"
message += f"• Eligible: {'✅ Yes' if stats['optimization_eligible'] else '❌ No'}\n"
message += f"• Reason: {stats['optimization_reason']}\n"
message += f"• New entries since last optimization: {stats['entries_since_last_optimization']}"
return message
except Exception as e:
return f"Error getting memory stats: {str(e)}"
@app.tool(
name="configure_memory_optimization",
description="Configure memory optimization settings for auto-optimization behavior.",
tags={"public", "memory"},
annotations={
"idempotentHint": False,
"readOnlyHint": False,
"title": "Configure Memory Optimization",
"parameters": {
"memory_file": "Optional path to specific memory file. If not provided, will configure the user's main memory file.",
"auto_optimize": "Enable or disable automatic optimization. True/False.",
"token_growth_threshold": "Token growth multiplier for triggering optimization (e.g., 1.20 = 20% growth). Must be >= 1.0.",
},
"returns": "Returns confirmation of updated settings.",
},
meta={
"category": "memory",
},
)
def configure_memory_optimization(
memory_file: Annotated[Optional[str], "Path to memory file to configure"] = None,
auto_optimize: Annotated[Optional[bool], "Enable/disable auto-optimization"] = None,
token_growth_threshold: Annotated[Optional[float], "Token growth threshold (e.g., 1.20 = 20% growth)"] = None,
) -> str:
"""Configure memory optimization settings."""
if read_only:
return "Error: Server is running in read-only mode"
try:
# Determine which file to configure
if memory_file:
file_path = Path(memory_file)
if not file_path.exists():
return f"Error: Memory file not found: {memory_file}"
else:
# Use default user memory file
user_memory_path = instruction_manager.get_memory_file_path()
if not user_memory_path.exists():
return "Error: No user memory file found to configure"
file_path = user_memory_path
# Read current frontmatter
from ..simple_file_ops import parse_frontmatter_file, write_frontmatter_file
frontmatter, content = parse_frontmatter_file(file_path)
# Update settings
updated_settings = []
if auto_optimize is not None:
frontmatter["autoOptimize"] = auto_optimize
updated_settings.append(f"auto_optimize: {auto_optimize}")
if token_growth_threshold is not None:
if token_growth_threshold < 1.0:
return "Error: token_growth_threshold must be >= 1.0 (e.g., 1.20 = 20% growth)"
frontmatter["tokenGrowthThreshold"] = token_growth_threshold
growth_percent = int((token_growth_threshold - 1.0) * 100)
updated_settings.append(f"token_growth_threshold: {token_growth_threshold} ({growth_percent}% growth)")
# Remove deprecated fields if they exist
deprecated_removed = []
for old_field in ["sizeThreshold", "entryThreshold", "timeThreshold"]:
if old_field in frontmatter:
frontmatter.pop(old_field)
deprecated_removed.append(old_field)
if not updated_settings and not deprecated_removed:
return "No settings provided to update. Available options: auto_optimize, token_growth_threshold"
# Write updated frontmatter
success = write_frontmatter_file(file_path, frontmatter, content, create_backup=True)
if success:
message = "✅ Memory optimization settings updated:\n"
for setting in updated_settings:
message += f"• {setting}\n"
if deprecated_removed:
message += "\n🧹 Removed deprecated settings: " + ", ".join(deprecated_removed) + "\n"
message += "\n💾 Backup created for safety"
return message
else:
return "❌ Failed to update memory optimization settings"
except Exception as e:
return f"Error configuring memory optimization: {str(e)}"
