atlas/home-voice-agent/mcp-server/tools/memory_tools.py

"""
MCP tools for memory management.

Allows LLM to read and write to long-term memory.
"""

import sys
from pathlib import Path

# Add parent directories to path
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from typing import Dict, Any
from tools.base import BaseTool
from memory.manager import get_memory_manager
from memory.schema import MemoryCategory, MemorySource

logger = None
try:
    import logging
    logger = logging.getLogger(__name__)
except:
    pass


class StoreMemoryTool(BaseTool):
    """Store a fact in long-term memory."""
    
    def __init__(self):
        self._name = "store_memory"
        self._description = "Store a fact in long-term memory. Use this when the user explicitly states a fact about themselves, their preferences, or routines."
        self._parameters = {
            "type": "object",
            "properties": {
                "category": {
                    "type": "string",
                    "enum": ["personal", "family", "preferences", "routines", "facts"],
                    "description": "Category of the memory"
                },
                "key": {
                    "type": "string",
                    "description": "Key for the memory (e.g., 'favorite_color', 'morning_routine')"
                },
                "value": {
                    "type": "string",
                    "description": "Value of the memory (e.g., 'blue', 'coffee at 7am')"
                },
                "confidence": {
                    "type": "number",
                    "minimum": 0.0,
                    "maximum": 1.0,
                    "default": 1.0,
                    "description": "Confidence level (1.0 for explicit, 0.7-0.9 for inferred)"
                },
                "context": {
                    "type": "string",
                    "description": "Additional context about the memory"
                }
            },
            "required": ["category", "key", "value"]
        }
        self.memory_manager = get_memory_manager()
    
    @property
    def name(self) -> str:
        return self._name
    
    @property
    def description(self) -> str:
        return self._description
    
    def get_schema(self):
        return {
            "name": self._name,
            "description": self._description,
            "inputSchema": {
                "type": "object",
                "properties": self._parameters["properties"],
                "required": self._parameters.get("required", [])
            }
        }
    
    def execute(self, arguments: Dict[str, Any]):
        """Store a memory entry."""
        category_str = arguments.get("category")
        key = arguments.get("key")
        value = arguments.get("value")
        confidence = arguments.get("confidence", 1.0)
        context = arguments.get("context")
        
        # Map string to enum
        category_map = {
            "personal": MemoryCategory.PERSONAL,
            "family": MemoryCategory.FAMILY,
            "preferences": MemoryCategory.PREFERENCES,
            "routines": MemoryCategory.ROUTINES,
            "facts": MemoryCategory.FACTS
        }
        
        category = category_map.get(category_str)
        if not category:
            raise ValueError(f"Invalid category: {category_str}")
        
        # Determine source based on confidence
        if confidence >= 0.95:
            source = MemorySource.EXPLICIT
        elif confidence >= 0.7:
            source = MemorySource.INFERRED
        else:
            source = MemorySource.CONFIRMED  # User confirmed inferred fact
        
        # Store memory
        entry = self.memory_manager.store_fact(
            category=category,
            key=key,
            value=value,
            confidence=confidence,
            source=source,
            context=context
        )
        
        return {
            "success": True,
            "message": f"Stored memory: {category_str}/{key} = {value}",
            "entry_id": entry.id,
            "confidence": entry.confidence
        }


class GetMemoryTool(BaseTool):
    """Get a fact from long-term memory."""
    
    def __init__(self):
        self._name = "get_memory"
        self._description = "Get a fact from long-term memory by category and key."
        self._parameters = {
                "type": "object",
                "properties": {
                    "category": {
                        "type": "string",
                        "enum": ["personal", "family", "preferences", "routines", "facts"],
                        "description": "Category of the memory"
                    },
                    "key": {
                        "type": "string",
                        "description": "Key for the memory"
                    }
                },
                "required": ["category", "key"]
            }
        self.memory_manager = get_memory_manager()
    
    @property
    def name(self) -> str:
        return self._name
    
    @property
    def description(self) -> str:
        return self._description
    
    def get_schema(self):
        return {
            "name": self._name,
            "description": self._description,
            "inputSchema": {
                "type": "object",
                "properties": self._parameters["properties"],
                "required": self._parameters.get("required", [])
            }
        }
    
    def execute(self, arguments: Dict[str, Any]):
        """Get a memory entry."""
        category_str = arguments.get("category")
        key = arguments.get("key")
        
        # Map string to enum
        category_map = {
            "personal": MemoryCategory.PERSONAL,
            "family": MemoryCategory.FAMILY,
            "preferences": MemoryCategory.PREFERENCES,
            "routines": MemoryCategory.ROUTINES,
            "facts": MemoryCategory.FACTS
        }
        
        category = category_map.get(category_str)
        if not category:
            raise ValueError(f"Invalid category: {category_str}")
        
        # Get memory
        entry = self.memory_manager.get_fact(category, key)
        
        if entry:
            return {
                "found": True,
                "category": category_str,
                "key": entry.key,
                "value": entry.value,
                "confidence": entry.confidence,
                "source": entry.source.value,
                "context": entry.context
            }
        else:
            return {
                "found": False,
                "message": f"No memory found for {category_str}/{key}"
            }


class SearchMemoryTool(BaseTool):
    """Search memory entries by query."""
    
    def __init__(self):
        self._name = "search_memory"
        self._description = "Search memory entries by query string. Useful for finding related facts."
        self._parameters = {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "Search query"
                    },
                    "category": {
                        "type": "string",
                        "enum": ["personal", "family", "preferences", "routines", "facts"],
                        "description": "Optional category filter"
                    },
                    "limit": {
                        "type": "integer",
                        "default": 10,
                        "description": "Maximum number of results"
                    }
                },
                "required": ["query"]
            }
        self.memory_manager = get_memory_manager()
    
    @property
    def name(self) -> str:
        return self._name
    
    @property
    def description(self) -> str:
        return self._description
    
    def get_schema(self):
        return {
            "name": self._name,
            "description": self._description,
            "inputSchema": {
                "type": "object",
                "properties": self._parameters["properties"],
                "required": self._parameters.get("required", [])
            }
        }
    
    def execute(self, arguments: Dict[str, Any]):
        """Search memory entries."""
        query = arguments.get("query")
        category_str = arguments.get("category")
        limit = arguments.get("limit", 10)
        
        category = None
        if category_str:
            category_map = {
                "personal": MemoryCategory.PERSONAL,
                "family": MemoryCategory.FAMILY,
                "preferences": MemoryCategory.PREFERENCES,
                "routines": MemoryCategory.ROUTINES,
                "facts": MemoryCategory.FACTS
            }
            category = category_map.get(category_str)
            if not category:
                raise ValueError(f"Invalid category: {category_str}")
        
        # Search memory
        results = self.memory_manager.search_facts(query, category, limit)
        
        return {
            "query": query,
            "count": len(results),
            "results": [
                {
                    "category": entry.category.value,
                    "key": entry.key,
                    "value": entry.value,
                    "confidence": entry.confidence,
                    "source": entry.source.value
                }
                for entry in results
            ]
        }


class ListMemoryTool(BaseTool):
    """List all memory entries in a category."""
    
    def __init__(self):
        self._name = "list_memory"
        self._description = "List all memory entries in a category."
        self._parameters = {
                "type": "object",
                "properties": {
                    "category": {
                        "type": "string",
                        "enum": ["personal", "family", "preferences", "routines", "facts"],
                        "description": "Category to list"
                    },
                    "limit": {
                        "type": "integer",
                        "default": 20,
                        "description": "Maximum number of entries"
                    }
                },
                "required": ["category"]
            }
        self.memory_manager = get_memory_manager()
    
    @property
    def name(self) -> str:
        return self._name
    
    @property
    def description(self) -> str:
        return self._description
    
    def get_schema(self):
        return {
            "name": self._name,
            "description": self._description,
            "inputSchema": {
                "type": "object",
                "properties": self._parameters["properties"],
                "required": self._parameters.get("required", [])
            }
        }
    
    def execute(self, arguments: Dict[str, Any]):
        """List memory entries in category."""
        category_str = arguments.get("category")
        limit = arguments.get("limit", 20)
        
        category_map = {
            "personal": MemoryCategory.PERSONAL,
            "family": MemoryCategory.FAMILY,
            "preferences": MemoryCategory.PREFERENCES,
            "routines": MemoryCategory.ROUTINES,
            "facts": MemoryCategory.FACTS
        }
        
        category = category_map.get(category_str)
        if not category:
            raise ValueError(f"Invalid category: {category_str}")
        
        # Get category facts
        entries = self.memory_manager.get_category_facts(category, limit)
        
        return {
            "category": category_str,
            "count": len(entries),
            "entries": [
                {
                    "key": entry.key,
                    "value": entry.value,
                    "confidence": entry.confidence,
                    "source": entry.source.value
                }
                for entry in entries
            ]
        }