Add memory_server.py

memory_server.py

@@ -0,0 +1,202 @@
+"""Memory Server for MCP protocol.
+
+This server exposes tools for storing, retrieving, and managing key/value pairs
+with optional namespace support. The data is persisted to a JSON file with
+metadata (timestamp).
+"""
+
+import json
+import fnmatch
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Optional, List
+
+from fastmcp import FastMCP
+
+
+class MemoryServer:
+    def __init__(self):
+        self.mcp = FastMCP("Memory-Server")
+        self.storage_path = Path("./memory_data.json")
+        # Load existing data into memory for quick access
+        self._data = self._load_memory()
+
+    def _load_memory(self) -> dict:
+        """Load memory from JSON file.
+
+        Returns:
+            dict: Mapping of keys to value+timestamp dicts.
+        """
+        if not self.storage_path.exists():
+            return {}
+        with open(self.storage_path, "r", encoding="utf-8") as f:
+            return json.load(f)
+
+    def _save_memory(self, data: dict):
+        """Persist memory to JSON file.
+
+        Args:
+            data (dict): The memory mapping to write.
+        """
+        self.storage_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(self.storage_path, "w", encoding="utf-8") as f:
+            json.dump(data, f, indent=2, ensure_ascii=False)
+
+    # ---------- Basic tools ----------
+    @self.mcp.tool()
+    def save(self, key: str, value: Any) -> bool:
+        """Save a value under a key.
+
+        Args:
+            key: Unique identifier.
+            value: Any JSON‑serializable value.
+
+        Returns:
+            True on success.
+        """
+        try:
+            self._data[key] = {
+                "value": value,
+                "timestamp": datetime.utcnow().isoformat()
+            }
+            self._save_memory(self._data)
+            return True
+        except Exception as e:
+            print(f"Error saving key {key}: {e}")
+            return False
+
+    @self.mcp.tool()
+    def get(self, key: str) -> Optional[dict]:
+        """Retrieve a value by key.
+
+        Args:
+            key: Identifier.
+
+        Returns:
+            Dict with fields ``key``, ``value``, ``timestamp`` or None.
+        """
+        entry = self._data.get(key)
+        if entry is None:
+            return None
+        return {
+            "key": key,
+            "value": entry["value"],
+            "timestamp": entry["timestamp"],
+        }
+
+    @self.mcp.tool()
+    def delete(self, key: str) -> bool:
+        """Delete a key from memory.
+
+        Args:
+            key: Identifier.
+
+        Returns:
+            True if deleted, False if key not found.
+        """
+        if key in self._data:
+            del self._data[key]
+            self._save_memory(self._data)
+            return True
+        return False
+
+    @self.mcp.tool()
+    def list_keys(self, pattern: str = "*") -> List[str]:
+        """List all keys matching a wildcard pattern.
+
+        Args:
+            pattern: Wildcard pattern (supports ``*`` and ``?``).
+
+        Returns:
+            List of matching keys.
+        """
+        return [k for k in self._data.keys() if fnmatch.fnmatch(k, pattern)]
+
+    # ---------- Namespace tools ----------
+    @self.mcp.tool()
+    def save_with_namespace(self, key: str, value: Any, namespace: str = "default") -> bool:
+        """Save a value under a namespaced key.
+
+        Args:
+            key: Identifier.
+            value: Value to store.
+            namespace: Namespace prefix.
+
+        Returns:
+            True on success.
+        """
+        composite_key = f"{namespace}:{key}"
+        return self.save(composite_key, value)
+
+    @self.mcp.tool()
+    def get_by_namespace(self, namespace: str = "default") -> List[dict]:
+        """Retrieve all entries in a namespace.
+
+        Args:
+            namespace: Namespace to query.
+
+        Returns:
+            List of dicts with ``key``, ``value``, ``timestamp``.
+        """
+        prefix = f"{namespace}:"
+        results = []
+        for k, v in self._data.items():
+            if k.startswith(prefix):
+                results.append({
+                    "key": k,
+                    "value": v["value"],
+                    "timestamp": v["timestamp"],
+                })
+        return results
+
+    @self.mcp.tool()
+    def get_namespace_key(self, key: str, namespace: str = "default") -> Optional[dict]:
+        """Retrieve a single key within a namespace.
+
+        Args:
+            key: Identifier without namespace.
+            namespace: Namespace.
+
+        Returns:
+            Entry dict or None.
+        """
+        composite_key = f"{namespace}:{key}"
+        return self.get(composite_key)
+
+    @self.mcp.tool()
+    def delete_namespace_key(self, key: str, namespace: str = "default") -> bool:
+        """Delete a single key within a namespace.
+
+        Args:
+            key: Identifier without namespace.
+            namespace: Namespace.
+
+        Returns:
+            True if deleted.
+        """
+        composite_key = f"{namespace}:{key}"
+        return self.delete(composite_key)
+
+    @self.mcp.tool()
+    def list_namespace_keys(self, namespace: str = "default", pattern: str = "*") -> List[str]:
+        """List keys in a namespace matching a pattern.
+
+        Args:
+            namespace: Namespace.
+            pattern: Wildcard pattern.
+
+        Returns:
+            List of matching keys (including namespace prefix).
+        """
+        prefix = f"{namespace}:"
+        return [k for k in self._data.keys() if k.startswith(prefix) and fnmatch.fnmatch(k, pattern)]
+
+
+# ---------- Server entry point ----------
+if __name__ == "__main__":
+    server = MemoryServer()
+    server.mcp.run(
+        transport="stdio",
+        show_banner=False,
+        log_level="ERROR",
+    )