t6_mem0_v2/mcp_server/main.py

"""
T6 Mem0 v2 MCP Server
Model Context Protocol server for memory operations
"""

import logging
import sys
import asyncio
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import (
    Resource,
    Tool,
    TextContent,
    ImageContent,
    EmbeddedResource
)
from mem0 import Memory

from config import mem0_config, settings
from mcp_server.tools import MemoryTools

# Configure logging
logging.basicConfig(
    level=getattr(logging, settings.log_level.upper()),
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stderr)]  # MCP uses stderr for logs
)

logger = logging.getLogger(__name__)


class T6Mem0Server:
    """T6 Mem0 v2 MCP Server"""

    def __init__(self):
        """Initialize MCP server"""
        self.server = Server("t6-mem0-v2")
        self.memory: Memory | None = None
        self.tools: MemoryTools | None = None

        # Setup handlers
        self.setup_handlers()

    def setup_handlers(self):
        """Setup MCP server handlers"""

        @self.server.list_resources()
        async def list_resources() -> list[Resource]:
            """
            List available resources (for future extension)
            """
            return []

        @self.server.read_resource()
        async def read_resource(uri: str) -> str:
            """
            Read resource by URI (for future extension)
            """
            logger.warning(f"Resource read not implemented: {uri}")
            return ""

        @self.server.list_tools()
        async def list_tools() -> list[Tool]:
            """List available memory tools"""
            logger.info("Listing tools")
            if not self.tools:
                raise RuntimeError("Tools not initialized")
            return self.tools.get_tool_definitions()

        @self.server.call_tool()
        async def call_tool(name: str, arguments: dict) -> list[TextContent | ImageContent | EmbeddedResource]:
            """
            Handle tool calls

            Args:
                name: Tool name
                arguments: Tool arguments

            Returns:
                Tool response
            """
            logger.info(f"Tool called: {name}")
            logger.debug(f"Arguments: {arguments}")

            if not self.tools:
                raise RuntimeError("Tools not initialized")

            # Route to appropriate handler
            handlers = {
                "add_memory": self.tools.handle_add_memory,
                "search_memories": self.tools.handle_search_memories,
                "get_memory": self.tools.handle_get_memory,
                "get_all_memories": self.tools.handle_get_all_memories,
                "update_memory": self.tools.handle_update_memory,
                "delete_memory": self.tools.handle_delete_memory,
                "delete_all_memories": self.tools.handle_delete_all_memories,
            }

            handler = handlers.get(name)
            if not handler:
                logger.error(f"Unknown tool: {name}")
                return [TextContent(type="text", text=f"Error: Unknown tool '{name}'")]

            try:
                return await handler(arguments)
            except Exception as e:
                logger.error(f"Tool execution failed: {e}", exc_info=True)
                return [TextContent(type="text", text=f"Error executing tool: {str(e)}")]

    async def initialize(self):
        """Initialize memory service"""
        logger.info("Initializing T6 Mem0 v2 MCP Server")
        logger.info(f"Environment: {settings.environment}")

        try:
            # Initialize Mem0
            logger.info("Initializing Mem0...")
            self.memory = Memory.from_config(config_dict=mem0_config)
            logger.info("Mem0 initialized successfully")

            # Initialize tools
            self.tools = MemoryTools(self.memory)
            logger.info("Tools initialized successfully")

            logger.info("T6 Mem0 v2 MCP Server ready")

        except Exception as e:
            logger.error(f"Failed to initialize server: {e}", exc_info=True)
            raise

    async def run(self):
        """Run the MCP server"""
        try:
            # Initialize before running
            await self.initialize()

            # Run server with stdio transport
            logger.info("Starting MCP server with stdio transport")
            async with stdio_server() as (read_stream, write_stream):
                await self.server.run(
                    read_stream,
                    write_stream,
                    self.server.create_initialization_options()
                )

        except Exception as e:
            logger.error(f"Server error: {e}", exc_info=True)
            raise


async def main():
    """Main entry point"""
    try:
        server = T6Mem0Server()
        await server.run()
    except KeyboardInterrupt:
        logger.info("Server stopped by user")
    except Exception as e:
        logger.error(f"Fatal error: {e}", exc_info=True)
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())