python-sdk/pkg/hanzo-tools-database/hanzo_tools/database/memory_tool.py

"""Unified memory tool with hybrid storage backend.

Combines plaintext markdown files, SQLite FTS, and vector search.
"""

import os
from typing import (
    Any,
    Dict,
    Unpack,
    Optional,
    Annotated,
    TypedDict,
    final,
    override,
)

from pydantic import Field
from mcp.server.fastmcp import Context as MCPContext

from hanzo_tools.core import BaseTool, PermissionManager, auto_timeout

from .memory_manager import MemoryManager

# Parameter types
Action = Annotated[
    str,
    Field(
        description="Memory action: read, write, append, search, create, list, stats",
        default="search",
    ),
]

Content = Annotated[
    Optional[str],
    Field(
        description="Content to store or search query",
        default=None,
    ),
]

FilePath = Annotated[
    Optional[str],
    Field(
        description="Markdown file path (e.g., 'rules.md', 'sessions/today.md')",
        default=None,
    ),
]

Category = Annotated[
    Optional[str],
    Field(
        description="Memory category for organization",
        default=None,
    ),
]

Scope = Annotated[
    str,
    Field(
        description="Memory scope: global, project, or both",
        default="project",
    ),
]

SearchType = Annotated[
    str,
    Field(
        description="Search type: fulltext, vector, or hybrid",
        default="fulltext",
    ),
]

Importance = Annotated[
    int,
    Field(
        description="Memory importance (1-10)",
        default=5,
    ),
]

Limit = Annotated[
    int,
    Field(
        description="Maximum results to return",
        default=10,
    ),
]


class MemoryParams(TypedDict, total=False):
    """Parameters for memory tool."""

    action: str
    content: Optional[str]
    file_path: Optional[str]
    category: Optional[str]
    scope: str
    search_type: str
    importance: int
    limit: int


@final
class MemoryTool(BaseTool):
    """Unified memory tool with hybrid storage."""

    def __init__(self, permission_manager: PermissionManager):
        """Initialize the memory tool."""
        super().__init__(permission_manager)
        self.memory_manager = MemoryManager(permission_manager)

    @property
    @override
    def name(self) -> str:
        """Get the tool name."""
        return "memory"

    @property
    @override
    def description(self) -> str:
        """Get the tool description."""
        return """Hybrid memory management. Actions: read, write, append, search, create, list, stats.

Storage types:
- Markdown files: Human-readable rule/context files
- SQLite FTS: Full-text search across all content
- Vector search: Semantic similarity (when available)

Usage:
memory --action read --file-path rules.md --scope global
memory --action write --file-path architecture.md --content "..." 
memory --action append --file-path sessions/today.md --content "New insight"
memory --action search --content "database design" --scope both
memory --action create --content "Important fact" --category general
memory --action list --scope project
memory --action stats
"""

    @override
    @auto_timeout("memory")
    async def call(
        self,
        ctx: MCPContext,
        **params: Unpack[MemoryParams],
    ) -> str:
        """Execute memory operation."""
        tool_ctx = self.create_tool_context(ctx)

        # Extract action
        action = params.get("action", "search")

        # Get current working directory for project context
        project_path = os.getcwd()

        # Route to appropriate handler
        try:
            if action == "read":
                return await self._handle_read(params, project_path, tool_ctx)
            elif action == "write":
                return await self._handle_write(params, project_path, tool_ctx)
            elif action == "append":
                return await self._handle_append(params, project_path, tool_ctx)
            elif action == "search":
                return await self._handle_search(params, project_path, tool_ctx)
            elif action == "create":
                return await self._handle_create(params, project_path, tool_ctx)
            elif action == "list":
                return await self._handle_list(params, project_path, tool_ctx)
            elif action == "stats":
                return await self._handle_stats(params, project_path, tool_ctx)
            else:
                return (
                    f"Error: Unknown action '{action}'. Valid actions: read, write, append, search, create, list, stats"
                )

        except Exception as e:
            await tool_ctx.error(f"Memory operation failed: {str(e)}")
            return f"Error: {str(e)}"

    async def _handle_read(self, params: Dict[str, Any], project_path: str, tool_ctx) -> str:
        """Read markdown memory file."""
        file_path = params.get("file_path")
        if not file_path:
            return "Error: file_path required for read action"

        scope = params.get("scope", "project")

        content = self.memory_manager.read_markdown_file(file_path, scope, project_path)

        if content is None:
            return f"Error: File '{file_path}' not found in {scope} scope"

        await tool_ctx.info(f"Read {scope} memory file: {file_path}")

        return f"=== {file_path} ({scope}) ===\n\n{content}"

    async def _handle_write(self, params: Dict[str, Any], project_path: str, tool_ctx) -> str:
        """Write markdown memory file."""
        file_path = params.get("file_path")
        content = params.get("content")

        if not file_path:
            return "Error: file_path required for write action"
        if not content:
            return "Error: content required for write action"

        scope = params.get("scope", "project")
        category = params.get("category")

        success = self.memory_manager.write_markdown_file(file_path, content, scope, project_path, category)

        if success:
            await tool_ctx.info(f"Wrote {scope} memory file: {file_path}")
            return f"Successfully wrote {scope} memory file: {file_path}"
        else:
            return f"Error: Failed to write memory file: {file_path}"

    async def _handle_append(self, params: Dict[str, Any], project_path: str, tool_ctx) -> str:
        """Append to markdown memory file."""
        file_path = params.get("file_path")
        content = params.get("content")

        if not file_path:
            return "Error: file_path required for append action"
        if not content:
            return "Error: content required for append action"

        scope = params.get("scope", "project")
        category = params.get("category")

        success = self.memory_manager.append_markdown_file(file_path, content, scope, project_path, category)

        if success:
            await tool_ctx.info(f"Appended to {scope} memory file: {file_path}")
            return f"Successfully appended to {scope} memory file: {file_path}"
        else:
            return f"Error: Failed to append to memory file: {file_path}"

    async def _handle_search(self, params: Dict[str, Any], project_path: str, tool_ctx) -> str:
        """Search across all memory types."""
        query = params.get("content")
        if not query:
            return "Error: content (search query) required for search action"

        scope = params.get("scope", "both")
        search_type = params.get("search_type", "fulltext")
        limit = params.get("limit", 10)

        results = self.memory_manager.search_memories(query, scope, project_path, search_type, limit)

        if not results:
            return f"No memories found for query: '{query}'"

        # Format results
        output = [f"=== Memory Search Results for '{query}' ==="]
        output.append(f"Found {len(results)} results (scope: {scope})\n")

        for i, result in enumerate(results, 1):
            source = result["source"]
            scope_name = result["scope"]

            if source == "markdown":
                output.append(f"{i}. [{scope_name}] {result['path']} (markdown)")
                if "snippet" in result:
                    output.append(f"   {result['snippet']}")
                output.append(f"   Category: {result.get('category', 'none')}")
                output.append(f"   Modified: {result['modified_at']}")

            elif source == "memory":
                output.append(f"{i}. [{scope_name}] Memory #{result['id']} (structured)")
                if "snippet" in result:
                    output.append(f"   {result['snippet']}")
                output.append(f"   Category: {result.get('category', 'none')}")
                output.append(f"   Importance: {result['importance']}")
                output.append(f"   Created: {result['created_at']}")

            output.append("")

        await tool_ctx.info(f"Found {len(results)} memories for: {query}")

        return "\n".join(output)

    async def _handle_create(self, params: Dict[str, Any], project_path: str, tool_ctx) -> str:
        """Create structured memory record."""
        content = params.get("content")
        if not content:
            return "Error: content required for create action"

        scope = params.get("scope", "project")
        category = params.get("category")
        importance = params.get("importance", 5)

        memory_id = self.memory_manager.create_memory(content, category, importance, None, scope, project_path)

        await tool_ctx.info(f"Created {scope} memory: #{memory_id}")

        return f"Successfully created {scope} memory #{memory_id}\nContent: {content[:100]}..."

    async def _handle_list(self, params: Dict[str, Any], project_path: str, tool_ctx) -> str:
        """List memory files and records."""
        scope = params.get("scope", "both")

        # List markdown files
        markdown_files = self.memory_manager.list_markdown_files(scope, project_path)

        output = [f"=== Memory Files ({scope} scope) ==="]

        if markdown_files:
            output.append("\nMarkdown Files:")
            for file_info in markdown_files:
                size_kb = file_info["size"] / 1024
                output.append(
                    f"  [{file_info['scope']}] {file_info['path']} ({size_kb:.1f}KB, {file_info['modified']})"
                )
        else:
            output.append("\nNo markdown files found")

        # Get stats for structured memories
        stats = self.memory_manager.get_memory_stats(scope, project_path)

        output.append(f"\nStructured Memories:")
        if scope in ("global", "both"):
            output.append(f"  Global: {stats['structured_memories']['global']} records")
        if scope in ("project", "both"):
            output.append(f"  Project: {stats['structured_memories']['project']} records")

        total_size_mb = stats["total_size_bytes"] / 1024 / 1024
        output.append(f"\nTotal size: {total_size_mb:.2f}MB")

        return "\n".join(output)

    async def _handle_stats(self, params: Dict[str, Any], project_path: str, tool_ctx) -> str:
        """Get detailed memory statistics."""
        scope = params.get("scope", "both")

        stats = self.memory_manager.get_memory_stats(scope, project_path)

        output = [f"=== Memory System Statistics ({scope} scope) ==="]
        output.append(f"Project: {project_path}\n")

        # Markdown files
        output.append("Markdown Files:")
        output.append(f"  Global: {stats['markdown_files']['global']} files")
        output.append(f"  Project: {stats['markdown_files']['project']} files")

        # Structured memories
        output.append("\nStructured Memories:")
        output.append(f"  Global: {stats['structured_memories']['global']} records")
        output.append(f"  Project: {stats['structured_memories']['project']} records")

        # Vector embeddings
        output.append("\nVector Embeddings:")
        output.append(f"  Global: {stats['vector_embeddings']['global']} embeddings")
        output.append(f"  Project: {stats['vector_embeddings']['project']} embeddings")

        # Storage size
        total_size_mb = stats["total_size_bytes"] / 1024 / 1024
        output.append(f"\nTotal Storage: {total_size_mb:.2f}MB")

        # Features status
        output.append("\nFeature Status:")
        output.append("  ✓ Markdown files")
        output.append("  ✓ SQLite FTS5 full-text search")

        # Check sqlite-vec availability
        try:
            import sqlite3

            conn = sqlite3.connect(":memory:")
            conn.enable_load_extension(True)
            conn.load_extension("vec0")
            output.append("  ✓ sqlite-vec vector search")
            conn.close()
        except Exception:
            output.append("  ✗ sqlite-vec vector search (not available)")

        return "\n".join(output)

    def register(self, mcp_server) -> None:
        """Register this tool with the MCP server."""
        pass