100 lines
3.2 KiB
Python
100 lines
3.2 KiB
Python
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()
|