main.py

#!/usr/bin/env -S uv run --script
# -*- coding: utf-8 -*-
# /// script
# dependencies = [
#     "mcp>=0.3.0",
#     "sqlite-vec>=0.1.6",
#     "sentence-transformers>=2.2.2"
# ]
# requires-python = ">=3.8"
# ///

"""
Vector Memory MCP Server - Main Entry Point
===========================================

A secure, vector-based memory server using sqlite-vec for semantic search.
Stores and retrieves coding memories, experiences, and knowledge using 
384-dimensional embeddings generated by sentence-transformers.

Usage:
    python main.py --working-dir /path/to/project

Memory files stored in: {working_dir}/memory/vector_memory.db
"""

import sys
from pathlib import Path
from typing import Dict, Any

# Add src to path for imports
sys.path.insert(0, str(Path(__file__).parent / "src"))

from mcp.server.fastmcp import FastMCP

# Import our modules
from src.models import Config
from src.security import validate_working_dir, SecurityError
from src.memory_store import VectorMemoryStore


def get_working_dir() -> Path:
    """Get working directory from command line arguments"""
    if len(sys.argv) >= 3 and sys.argv[1] == "--working-dir":
        return validate_working_dir(sys.argv[2])
    else:
        # Default to current directory
        return validate_working_dir(".")


def create_server() -> FastMCP:
    """Create and configure the MCP server"""
    
    # Initialize global memory store
    try:
        memory_dir = get_working_dir()
        db_path = memory_dir / Config.DB_NAME
        memory_store = VectorMemoryStore(db_path)
        print(f"Memory database initialized: {db_path}", file=sys.stderr)
    except Exception as e:
        print(f"Failed to initialize memory store: {e}", file=sys.stderr)
        sys.exit(1)
    
    # Create FastMCP server
    mcp = FastMCP(Config.SERVER_NAME)
    
    # ===============================================================================
    # MCP TOOLS IMPLEMENTATION
    # ===============================================================================
    
    @mcp.tool()
    def store_memory(
        content: str,
        category: str = "other",
        tags: list[str] = None
    ) -> dict[str, Any]:
        """
        Store a coding memory, experience, or knowledge with automatic vector embedding.
        
        Use this to remember:
        - Code solutions and patterns you've implemented
        - Bug fixes and debugging approaches
        - Architecture decisions and rationale
        - Tool usage and configuration tips
        - Performance optimizations
        - Security considerations
        - Learning insights
        
        Args:
            content: The memory content (max 10,000 characters)
            category: Memory category - one of: code-solution, bug-fix, architecture, 
                     learning, tool-usage, debugging, performance, security, other
            tags: List of relevant tags for organization (max 10 tags)
        
        Returns:
            Dict with success status, memory_id, and metadata
        """
        try:
            if tags is None:
                tags = []
            
            result = memory_store.store_memory(content, category, tags)
            return result
            
        except SecurityError as e:
            return {
                "success": False,
                "error": "Security validation failed",
                "message": str(e)
            }
        except Exception as e:
            return {
                "success": False,
                "error": "Storage failed",
                "message": str(e)
            }
    
    @mcp.tool()
    def search_memories(
        query: str,
        limit: int = 10,
        category: str = None
    ) -> dict[str, Any]:
        """
        Search stored memories using semantic similarity.
        
        This performs vector-based semantic search to find memories most relevant
        to your query, even if they don't contain exact keyword matches.
        
        Examples:
        - "authentication patterns" -> finds JWT, OAuth, login implementations
        - "async error handling" -> finds Promise, async/await, error catching
        - "database optimization" -> finds indexing, query performance, caching
        
        Args:
            query: Search query describing what you're looking for
            limit: Maximum number of results to return (1-50, default: 10)
            category: Optional category filter (code-solution, bug-fix, etc.)
        
        Returns:
            Dict with search results including content, similarity scores, and metadata
        """
        try:
            search_results = memory_store.search_memories(query, limit, category)
            
            if not search_results:
                return {
                    "success": True,
                    "results": [],
                    "message": "No matching memories found. Try different keywords or broader terms."
                }
            
            # Convert SearchResult objects to dictionaries
            results = [result.to_dict() for result in search_results]
            
            return {
                "success": True,
                "query": query,
                "results": results,
                "count": len(results),
                "message": f"Found {len(results)} relevant memories"
            }
            
        except SecurityError as e:
            return {
                "success": False,
                "error": "Security validation failed",
                "message": str(e)
            }
        except Exception as e:
            return {
                "success": False,
                "error": "Search failed",
                "message": str(e)
            }
    
    @mcp.tool()
    def list_recent_memories(limit: int = 10) -> dict[str, Any]:
        """
        List recently stored memories in chronological order.
        
        Useful for reviewing what you've recently learned or stored,
        and for getting familiar with the type of content in your memory.
        
        Args:
            limit: Maximum number of recent memories to return (1-50, default: 10)
        
        Returns:
            Dict with list of recent memories and their metadata
        """
        try:
            limit = min(max(1, limit), Config.MAX_MEMORIES_PER_SEARCH)
            memories = memory_store.get_recent_memories(limit)
            
            # Convert MemoryEntry objects to dictionaries
            memory_dicts = [memory.to_dict() for memory in memories]
            
            return {
                "success": True,
                "memories": memory_dicts,
                "count": len(memory_dicts),
                "message": f"Retrieved {len(memory_dicts)} recent memories"
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": "Failed to get recent memories",
                "message": str(e)
            }
    
    @mcp.tool()
    def get_memory_stats() -> dict[str, Any]:
        """
        Get comprehensive statistics about the memory database.
        
        Provides insights into:
        - Total number of stored memories
        - Memory usage and limits
        - Category breakdown
        - Recent activity
        - Most accessed memories
        - Database size and health
        
        Returns:
            Dict with detailed statistics and health information
        """
        try:
            stats = memory_store.get_stats()
            result = stats.to_dict()
            result["success"] = True
            return result
            
        except Exception as e:
            return {
                "success": False,
                "error": "Failed to get statistics",
                "message": str(e)
            }
    
    @mcp.tool()
    def clear_old_memories(
        days_old: int = 30,
        max_to_keep: int = 1000
    ) -> dict[str, Any]:
        """
        Clear old, less frequently accessed memories to free space.
        
        This tool helps maintain database performance by removing older memories
        that haven't been accessed recently. It prioritizes keeping frequently
        accessed and recent memories.
        
        Args:
            days_old: Only consider memories older than this many days (default: 30)
            max_to_keep: Maximum total memories to keep in database (default: 1000)
        
        Returns:
            Dict with cleanup results and statistics
        """
        try:
            if days_old < 1:
                return {
                    "success": False,
                    "error": "Invalid parameter",
                    "message": "days_old must be at least 1"
                }
            
            result = memory_store.clear_old_memories(days_old, max_to_keep)
            return result
            
        except SecurityError as e:
            return {
                "success": False,
                "error": "Security validation failed",
                "message": str(e)
            }
        except Exception as e:
            return {
                "success": False,
                "error": "Cleanup failed",
                "message": str(e)
            }
    
    return mcp


def main():
    """Main entry point"""
    print(f"Starting {Config.SERVER_NAME} v{Config.SERVER_VERSION}", file=sys.stderr)
    
    try:
        # Get working directory info
        memory_dir = get_working_dir()
        db_path = memory_dir / Config.DB_NAME
        
        print(f"Working directory: {memory_dir.parent}", file=sys.stderr)
        print(f"Memory database: {db_path}", file=sys.stderr)
        print(f"Embedding model: {Config.EMBEDDING_MODEL}", file=sys.stderr)
        print("=" * 50, file=sys.stderr)
        
        # Create and run server
        server = create_server()
        print("Server ready for connections...", file=sys.stderr)
        server.run()
        
    except KeyboardInterrupt:
        print("\nServer stopped by user", file=sys.stderr)
    except Exception as e:
        print(f"Server failed to start: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()