add memory_server.py
This commit is contained in:
+100
-139
@@ -1,178 +1,139 @@
|
||||
"""
|
||||
Memory Server for MCP protocol.
|
||||
Memory Server implementing MCP protocol using FastMCP.
|
||||
|
||||
Provides basic memory operations with optional namespace support.
|
||||
"""
|
||||
import json
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from fastmcp import FastMCP
|
||||
import fnmatch
|
||||
|
||||
class MemoryServer:
|
||||
"""FastMCP based memory server.
|
||||
"""A simple MCP server that stores key/value pairs in a JSON file.
|
||||
|
||||
Stores data in a JSON file located at ``./memory_data.json``.
|
||||
Each entry is stored as ``{namespace:key: {value, timestamp}}``.
|
||||
The server exposes six tools:
|
||||
* save – store a value under a key.
|
||||
* get – retrieve a stored value.
|
||||
* delete – remove a key.
|
||||
* list_keys – list all keys with optional glob pattern.
|
||||
* save_with_namespace – same as ``save`` but prefixes the key with a namespace.
|
||||
* get_by_namespace – return all items belonging to a namespace.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
def __init__(self, storage_path: str | Path = "./memory_data.json"):
|
||||
self.mcp = FastMCP("Memory-Server")
|
||||
self.storage_path = Path("./memory_data.json")
|
||||
|
||||
# ---------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ---------------------------------------------------------------------
|
||||
def _load_memory(self) -> Dict[str, Any]:
|
||||
if not self.storage_path.exists():
|
||||
return {}
|
||||
with self.storage_path.open("r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
def _save_memory(self, data: Dict[str, Any]) -> None:
|
||||
self.storage_path = Path(storage_path)
|
||||
# Ensure directory exists for future writes
|
||||
self.storage_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with self.storage_path.open("w", encoding="utf-8") as f:
|
||||
|
||||
def _load(self) -> Dict[str, Any]:
|
||||
"""Load the entire memory store from disk.
|
||||
|
||||
Returns an empty dict if file does not exist or is invalid.
|
||||
"""
|
||||
try:
|
||||
if self.storage_path.exists():
|
||||
with self.storage_path.open("r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
except Exception:
|
||||
# Corrupted file – start fresh
|
||||
pass
|
||||
return {}
|
||||
|
||||
def _save(self, data: Dict[str, Any]) -> None:
|
||||
"""Persist the memory store to disk.
|
||||
|
||||
The file is written atomically by writing to a temp file first.
|
||||
"""
|
||||
tmp = self.storage_path.with_suffix(".tmp")
|
||||
with tmp.open("w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2, ensure_ascii=False)
|
||||
tmp.replace(self.storage_path)
|
||||
|
||||
def _make_key(self, key: str, namespace: str = "default") -> str:
|
||||
return f"{namespace}:{key}"
|
||||
def _current_timestamp(self) -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
# ---------------------------------------------------------------------
|
||||
# Tools exposed via FastMCP
|
||||
# ---------------------------------------------------------------------
|
||||
@self.mcp.tool()
|
||||
def save(self, key: str, value: Any) -> bool:
|
||||
"""Save a value under ``default`` namespace.
|
||||
@property
|
||||
def mcp_tool(self):
|
||||
# Helper to expose the FastMCP instance for decorators
|
||||
return self.mcp
|
||||
|
||||
Args:
|
||||
key: Identifier for the value.
|
||||
value: Serializable data.
|
||||
# ---------- Basic tools ----------
|
||||
@FastMCP.tool()
|
||||
def save(key: str, value: Any) -> bool:
|
||||
"""Store *value* under *key*.
|
||||
|
||||
Returns:
|
||||
True if saved successfully.
|
||||
The value is JSON‑serialisable. Existing keys are overwritten.
|
||||
Returns ``True`` on success.
|
||||
"""
|
||||
return self.save_with_namespace(key, value, "default")
|
||||
|
||||
@self.mcp.tool()
|
||||
def get(self, key: str) -> Optional[Dict[str, Any]]:
|
||||
"""Retrieve a value from ``default`` namespace.
|
||||
|
||||
Args:
|
||||
key: Identifier to fetch.
|
||||
|
||||
Returns:
|
||||
Dictionary with ``value`` and ``timestamp`` or None if missing.
|
||||
"""
|
||||
return self.get_by_namespace("default").get(key)
|
||||
|
||||
@self.mcp.tool()
|
||||
def delete(self, key: str) -> bool:
|
||||
"""Delete a key from ``default`` namespace.
|
||||
|
||||
Args:
|
||||
key: Identifier to remove.
|
||||
|
||||
Returns:
|
||||
True if removed, False otherwise.
|
||||
"""
|
||||
return self.delete_from_namespace(key, "default")
|
||||
|
||||
@self.mcp.tool()
|
||||
def list_keys(self, pattern: str = "*") -> List[str]:
|
||||
"""List all keys in ``default`` namespace matching a glob pattern.
|
||||
|
||||
Args:
|
||||
pattern: Glob pattern (supports * and ?).
|
||||
|
||||
Returns:
|
||||
List of key names without namespace prefix.
|
||||
"""
|
||||
return self.list_keys_in_namespace("default", pattern)
|
||||
|
||||
@self.mcp.tool()
|
||||
def save_with_namespace(self, key: str, value: Any, namespace: str = "default") -> bool:
|
||||
"""Save a value under a specific namespace.
|
||||
|
||||
Args:
|
||||
key: Identifier for the value.
|
||||
value: Serializable data.
|
||||
namespace: Namespace name.
|
||||
|
||||
Returns:
|
||||
True if saved successfully.
|
||||
"""
|
||||
full_key = self._make_key(key, namespace)
|
||||
data = self._load_memory()
|
||||
data[full_key] = {
|
||||
"value": value,
|
||||
"timestamp": datetime.utcnow().isoformat() + "Z",
|
||||
}
|
||||
self._save_memory(data)
|
||||
data = self._load()
|
||||
data[key] = {"value": value, "timestamp": self._current_timestamp()}
|
||||
self._save(data)
|
||||
return True
|
||||
|
||||
@self.mcp.tool()
|
||||
def get_by_namespace(self, namespace: str = "default") -> Dict[str, Any]:
|
||||
"""Return all key/value pairs in a namespace.
|
||||
@FastMCP.tool()
|
||||
def get(key: str) -> Optional[Dict[str, Any]]:
|
||||
"""Retrieve the stored item for *key*.
|
||||
|
||||
Args:
|
||||
namespace: Namespace name.
|
||||
|
||||
Returns:
|
||||
Mapping of key to dict with ``value`` and ``timestamp``.
|
||||
Returns a dict with ``value`` and ``timestamp`` or ``None`` if not found.
|
||||
"""
|
||||
data = self._load_memory()
|
||||
prefix = f"{namespace}:"
|
||||
result: Dict[str, Any] = {}
|
||||
for k, v in data.items():
|
||||
if k.startswith(prefix):
|
||||
key_name = k[len(prefix) :]
|
||||
result[key_name] = v
|
||||
return result
|
||||
data = self._load()
|
||||
return data.get(key)
|
||||
|
||||
@self.mcp.tool()
|
||||
def delete_from_namespace(self, key: str, namespace: str = "default") -> bool:
|
||||
"""Delete a key from a specific namespace.
|
||||
@FastMCP.tool()
|
||||
def delete(key: str) -> bool:
|
||||
"""Remove *key* from the store.
|
||||
|
||||
Args:
|
||||
key: Identifier to remove.
|
||||
namespace: Namespace name.
|
||||
|
||||
Returns:
|
||||
True if removed, False otherwise.
|
||||
Returns ``True`` if key existed and was removed, otherwise ``False``.
|
||||
"""
|
||||
full_key = self._make_key(key, namespace)
|
||||
data = self._load_memory()
|
||||
if full_key in data:
|
||||
del data[full_key]
|
||||
self._save_memory(data)
|
||||
data = self._load()
|
||||
if key in data:
|
||||
del data[key]
|
||||
self._save(data)
|
||||
return True
|
||||
return False
|
||||
|
||||
@self.mcp.tool()
|
||||
def list_keys_in_namespace(self, namespace: str = "default", pattern: str = "*") -> List[str]:
|
||||
"""List keys in a namespace matching a glob pattern.
|
||||
@FastMCP.tool()
|
||||
def list_keys(pattern: str = "*") -> List[str]:
|
||||
"""Return all keys matching *pattern*.
|
||||
|
||||
Args:
|
||||
namespace: Namespace name.
|
||||
pattern: Glob pattern.
|
||||
|
||||
Returns:
|
||||
List of key names without namespace prefix.
|
||||
The pattern supports Unix shell wildcards (``*``, ``?``).
|
||||
"""
|
||||
import fnmatch
|
||||
data = self._load()
|
||||
return [k for k in data.keys() if fnmatch.fnmatch(k, pattern)]
|
||||
|
||||
data = self._load_memory()
|
||||
# ---------- Namespace tools ----------
|
||||
@FastMCP.tool()
|
||||
def save_with_namespace(key: str, value: Any, namespace: str = "default") -> bool:
|
||||
"""Store *value* under ``namespace:key``.
|
||||
|
||||
The key is prefixed with the namespace and a colon. Existing entries are overwritten.
|
||||
Returns ``True`` on success.
|
||||
"""
|
||||
namespaced_key = f"{namespace}:{key}"
|
||||
return self.save(namespaced_key, value)
|
||||
|
||||
@FastMCP.tool()
|
||||
def get_by_namespace(namespace: str = "default") -> List[Dict[str, Any]]:
|
||||
"""Return all items belonging to *namespace*.
|
||||
|
||||
Each returned dict contains ``key`` (without namespace prefix), ``value`` and ``timestamp``.
|
||||
"""
|
||||
data = self._load()
|
||||
prefix = f"{namespace}:"
|
||||
keys = [k[len(prefix) :] for k in data if k.startswith(prefix)]
|
||||
return [k for k in keys if fnmatch.fnmatch(k, pattern)]
|
||||
|
||||
# ---------------------------------------------------------------------
|
||||
# Entry point
|
||||
# ---------------------------------------------------------------------
|
||||
def run(self) -> None:
|
||||
self.mcp.run(transport="stdio", show_banner=False, log_level="ERROR")
|
||||
result: List[Dict[str, Any]] = []
|
||||
for full_key, meta in data.items():
|
||||
if full_key.startswith(prefix):
|
||||
key_without_ns = full_key[len(prefix) :]
|
||||
item = {"key": key_without_ns, "value": meta["value"], "timestamp": meta["timestamp"]}
|
||||
result.append(item)
|
||||
return result
|
||||
|
||||
# ---------- Server entry point ----------
|
||||
if __name__ == "__main__":
|
||||
server = MemoryServer()
|
||||
server.run()
|
||||
# Run with stdio transport – suitable for local testing and client usage.
|
||||
server.mcp.run(transport="stdio", show_banner=False, log_level="ERROR")
|
||||
|
||||
Reference in New Issue
Block a user