Update mcp_server.py with smart injection and real embeddings

mcp_server.py

@@ -0,0 +1,345 @@
+#!/usr/bin/env python3
+"""
+Mnemo MCP Server
+Provides memory capabilities to Claude and other MCP-compatible LLMs
+
+Usage:
+    uvx mnemo-memory  # Run as MCP server
+    
+Or in Claude config:
+    {
+        "mcpServers": {
+            "mnemo": {"command": "uvx", "args": ["mnemo-memory"]}
+        }
+    }
+"""
+
+import json
+import sys
+from typing import Any, Dict, List
+
+try:
+    from mcp.server import Server
+    from mcp.types import Tool, TextContent
+    HAS_MCP = True
+except ImportError:
+    HAS_MCP = False
+    print("Warning: MCP not installed. Install with: pip install mcp", file=sys.stderr)
+
+from mnemo import Mnemo, should_inject_memory
+
+# Global memory instance
+memory = Mnemo(use_real_embeddings=True)
+
+# Tool definitions
+TOOLS = [
+    {
+        "name": "mnemo_add",
+        "description": "Store a new memory. Use this to save important information for later retrieval.",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "content": {
+                    "type": "string",
+                    "description": "The content to remember"
+                },
+                "metadata": {
+                    "type": "object",
+                    "description": "Optional metadata (tags, source, etc.)"
+                }
+            },
+            "required": ["content"]
+        }
+    },
+    {
+        "name": "mnemo_search",
+        "description": "Search stored memories. Returns relevant memories ranked by similarity.",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "Search query"
+                },
+                "top_k": {
+                    "type": "integer",
+                    "description": "Number of results (default: 5)",
+                    "default": 5
+                }
+            },
+            "required": ["query"]
+        }
+    },
+    {
+        "name": "mnemo_should_inject",
+        "description": "Check if memory should be injected for this query. Uses context-check algorithm with 90% accuracy.",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "The query to check"
+                },
+                "context": {
+                    "type": "string",
+                    "description": "Optional additional context"
+                }
+            },
+            "required": ["query"]
+        }
+    },
+    {
+        "name": "mnemo_get_context",
+        "description": "Get formatted memory context for injection into prompts.",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "Search query for finding relevant context"
+                },
+                "top_k": {
+                    "type": "integer",
+                    "description": "Number of memories to include (default: 3)",
+                    "default": 3
+                }
+            },
+            "required": ["query"]
+        }
+    },
+    {
+        "name": "mnemo_list",
+        "description": "List all stored memories.",
+        "inputSchema": {
+            "type": "object",
+            "properties": {}
+        }
+    },
+    {
+        "name": "mnemo_delete",
+        "description": "Delete a specific memory by ID.",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "memory_id": {
+                    "type": "string",
+                    "description": "ID of the memory to delete"
+                }
+            },
+            "required": ["memory_id"]
+        }
+    },
+    {
+        "name": "mnemo_stats",
+        "description": "Get memory system statistics.",
+        "inputSchema": {
+            "type": "object",
+            "properties": {}
+        }
+    },
+    {
+        "name": "mnemo_clear",
+        "description": "Clear all stored memories. Use with caution!",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "confirm": {
+                    "type": "boolean",
+                    "description": "Must be true to confirm deletion"
+                }
+            },
+            "required": ["confirm"]
+        }
+    }
+]
+
+
+def handle_tool_call(name: str, arguments: Dict[str, Any]) -> str:
+    """Handle a tool call and return result"""
+    
+    if name == "mnemo_add":
+        content = arguments.get("content", "")
+        metadata = arguments.get("metadata", {})
+        mem_id = memory.add(content, metadata)
+        return json.dumps({
+            "status": "success",
+            "memory_id": mem_id,
+            "message": f"Memory stored successfully"
+        })
+    
+    elif name == "mnemo_search":
+        query = arguments.get("query", "")
+        top_k = arguments.get("top_k", 5)
+        results = memory.search(query, top_k=top_k)
+        
+        return json.dumps({
+            "status": "success",
+            "count": len(results),
+            "results": [
+                {
+                    "id": r.id,
+                    "content": r.content,
+                    "score": round(r.score, 3),
+                    "metadata": r.metadata
+                }
+                for r in results
+            ]
+        })
+    
+    elif name == "mnemo_should_inject":
+        query = arguments.get("query", "")
+        context = arguments.get("context", "")
+        should, reason = should_inject_memory(query, context)
+        
+        return json.dumps({
+            "should_inject": should,
+            "reason": reason,
+            "recommendation": "Inject memory context" if should else "Skip memory - standalone query"
+        })
+    
+    elif name == "mnemo_get_context":
+        query = arguments.get("query", "")
+        top_k = arguments.get("top_k", 3)
+        context = memory.get_context(query, top_k=top_k)
+        
+        return json.dumps({
+            "status": "success",
+            "context": context if context else None,
+            "message": "Context retrieved" if context else "No relevant context found"
+        })
+    
+    elif name == "mnemo_list":
+        memories = memory.list_all()
+        return json.dumps({
+            "status": "success",
+            "count": len(memories),
+            "memories": [
+                {
+                    "id": m.id,
+                    "content": m.content[:200] + "..." if len(m.content) > 200 else m.content,
+                    "created_at": m.created_at,
+                    "metadata": m.metadata
+                }
+                for m in memories
+            ]
+        })
+    
+    elif name == "mnemo_delete":
+        memory_id = arguments.get("memory_id", "")
+        success = memory.delete(memory_id)
+        
+        return json.dumps({
+            "status": "success" if success else "error",
+            "message": f"Memory {memory_id} deleted" if success else f"Memory {memory_id} not found"
+        })
+    
+    elif name == "mnemo_stats":
+        stats = memory.get_stats()
+        return json.dumps({
+            "status": "success",
+            "stats": stats
+        })
+    
+    elif name == "mnemo_clear":
+        if arguments.get("confirm", False):
+            memory.clear()
+            return json.dumps({
+                "status": "success",
+                "message": "All memories cleared"
+            })
+        else:
+            return json.dumps({
+                "status": "error",
+                "message": "Must set confirm=true to clear all memories"
+            })
+    
+    else:
+        return json.dumps({
+            "status": "error",
+            "message": f"Unknown tool: {name}"
+        })
+
+
+def run_mcp_server():
+    """Run as MCP server"""
+    if not HAS_MCP:
+        print("Error: MCP not installed. Install with: pip install mcp", file=sys.stderr)
+        sys.exit(1)
+    
+    server = Server("mnemo-memory")
+    
+    @server.list_tools()
+    async def list_tools() -> List[Tool]:
+        return [
+            Tool(
+                name=tool["name"],
+                description=tool["description"],
+                inputSchema=tool["inputSchema"]
+            )
+            for tool in TOOLS
+        ]
+    
+    @server.call_tool()
+    async def call_tool(name: str, arguments: Dict[str, Any]) -> List[TextContent]:
+        result = handle_tool_call(name, arguments)
+        return [TextContent(type="text", text=result)]
+    
+    # Run server
+    import asyncio
+    asyncio.run(server.run())
+
+
+def run_cli():
+    """Simple CLI for testing"""
+    print("Mnemo Memory System - CLI Mode")
+    print("Commands: add, search, inject, context, list, stats, clear, quit")
+    print("-" * 50)
+    
+    while True:
+        try:
+            cmd = input("\n> ").strip().lower()
+            
+            if cmd == "quit":
+                break
+            elif cmd == "add":
+                content = input("Content: ")
+                result = handle_tool_call("mnemo_add", {"content": content})
+                print(result)
+            elif cmd == "search":
+                query = input("Query: ")
+                result = handle_tool_call("mnemo_search", {"query": query})
+                print(result)
+            elif cmd == "inject":
+                query = input("Query: ")
+                result = handle_tool_call("mnemo_should_inject", {"query": query})
+                print(result)
+            elif cmd == "context":
+                query = input("Query: ")
+                result = handle_tool_call("mnemo_get_context", {"query": query})
+                print(result)
+            elif cmd == "list":
+                result = handle_tool_call("mnemo_list", {})
+                print(result)
+            elif cmd == "stats":
+                result = handle_tool_call("mnemo_stats", {})
+                print(result)
+            elif cmd == "clear":
+                confirm = input("Are you sure? (yes/no): ")
+                if confirm.lower() == "yes":
+                    result = handle_tool_call("mnemo_clear", {"confirm": True})
+                    print(result)
+            else:
+                print("Unknown command. Use: add, search, inject, context, list, stats, clear, quit")
+        
+        except KeyboardInterrupt:
+            print("\nBye!")
+            break
+        except Exception as e:
+            print(f"Error: {e}")
+
+
+if __name__ == "__main__":
+    if len(sys.argv) > 1 and sys.argv[1] == "--cli":
+        run_cli()
+    else:
+        run_mcp_server()