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