Add memory_server.py

memory_server.py

@@ -0,0 +1,99 @@
+import json
+import os
+import time
+from typing import Any, Dict, List, Optional
+
+from fastmcp import FastMCP, mcp
+from pydantic import BaseModel, ValidationError, Field
+
+# Path to store memory data
+DATA_FILE = os.path.join(os.path.dirname(__file__), "memory_data.json")
+
+
+class MemoryItem(BaseModel):
+    key: str
+    value: Any
+    timestamp: float
+
+
+class MemoryServer:
+    def __init__(self, data_file: str = DATA_FILE):
+        self.data_file = data_file
+        self.memory: Dict[str, MemoryItem] = {}
+        self._load_memory()
+        self.mcp = FastMCP()
+        self._register_tools()
+
+    def _load_memory(self):
+        if not os.path.exists(self.data_file):
+            self.memory = {}
+            return
+        try:
+            with open(self.data_file, "r", encoding="utf-8") as f:
+                raw = json.load(f)
+            for k, v in raw.items():
+                try:
+                    self.memory[k] = MemoryItem(**v)
+                except ValidationError:
+                    continue
+        except Exception:
+            self.memory = {}
+
+    def _save_memory(self):
+        try:
+            raw = {k: v.dict() for k, v in self.memory.items()}
+            with open(self.data_file, "w", encoding="utf-8") as f:
+                json.dump(raw, f, indent=2)
+        except Exception:
+            pass
+
+    def _register_tools(self):
+        @self.mcp.tool("save", "Saves a key-value pair.")
+        def save(key: str, value: Any) -> bool:
+            if not self._validate_key(key):
+                return False
+            self.memory[key] = MemoryItem(key=key, value=value, timestamp=time.time())
+            self._save_memory()
+            return True
+
+        @self.mcp.tool("get", "Retrieves a key-value pair.")
+        def get(key: str) -> Optional[Dict[str, Any]]:
+            item = self.memory.get(key)
+            if item:
+                return item.dict()
+            return None
+
+        @self.mcp.tool("delete", "Deletes a key-value pair.")
+        def delete(key: str) -> bool:
+            if key in self.memory:
+                del self.memory[key]
+                self._save_memory()
+                return True
+            return False
+
+        @self.mcp.tool("list_keys", "Lists keys matching pattern.")
+        def list_keys(pattern: str = "*") -> List[str]:
+            import fnmatch
+            return [k for k in self.memory if fnmatch.fnmatch(k, pattern)]
+
+        @self.mcp.tool("save_with_namespace", "Saves a key-value pair with namespace.")
+        def save_with_namespace(key: str, value: Any, namespace: str = "default") -> bool:
+            full_key = f"{namespace}:{key}"
+            return save(full_key, value)
+
+        @self.mcp.tool("get_by_namespace", "Gets all items in a namespace.")
+        def get_by_namespace(namespace: str = "default") -> List[Dict[str, Any]]:
+            prefix = f"{namespace}:"
+            return [item.dict() for k, item in self.memory.items() if k.startswith(prefix)]
+
+    def _validate_key(self, key: str) -> bool:
+        # disallow colon to avoid confusion with namespace separator
+        return ":" not in key
+
+    def run(self):
+        self.mcp.run()
+
+
+if __name__ == "__main__":
+    server = MemoryServer()
+    server.run()