diff --git a/memory_server.py b/memory_server.py new file mode 100644 index 0000000..f0ed0aa --- /dev/null +++ b/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()