235 lines
7.2 KiB
Python
235 lines
7.2 KiB
Python
#!/usr/bin/env python3
|
||
"""Main entry point for MCP memory server and client demonstration.
|
||
|
||
This file contains both the server implementation (MemoryServer) and a simple
|
||
client example that demonstrates how to use the server via FastMCP.
|
||
|
||
The server exposes six tools:
|
||
|
||
* save
|
||
* get
|
||
* delete
|
||
* list_keys
|
||
* save_with_namespace
|
||
* get_by_namespace
|
||
|
||
The client shows how to call these tools using the FastMCP Client.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import fnmatch
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Any, Optional, List, Dict
|
||
|
||
from fastmcp import FastMCP, Client
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Server implementation
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class MemoryServer:
|
||
"""Simple key‑value store backed by a JSON file.
|
||
|
||
Keys are stored as ``namespace:key``. Each entry contains the value and a
|
||
timestamp of the last write.
|
||
"""
|
||
|
||
def __init__(self, storage_path: str | Path = "./memory_data.json"):
|
||
self.mcp = FastMCP("Memory-Server")
|
||
self.storage_path = Path(storage_path)
|
||
self._ensure_storage()
|
||
|
||
# ---------------------------------------------------------------------
|
||
# Persistence helpers
|
||
# ---------------------------------------------------------------------
|
||
|
||
def _ensure_storage(self) -> None:
|
||
if not self.storage_path.parent.exists():
|
||
self.storage_path.parent.mkdir(parents=True, exist_ok=True)
|
||
if not self.storage_path.exists():
|
||
self.storage_path.write_text("{}", encoding="utf-8")
|
||
|
||
def _load(self) -> Dict[str, Dict[str, Any]]:
|
||
try:
|
||
return json.loads(self.storage_path.read_text(encoding="utf-8"))
|
||
except json.JSONDecodeError:
|
||
return {}
|
||
|
||
def _save(self, data: Dict[str, Dict[str, Any]]) -> None:
|
||
self.storage_path.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
|
||
|
||
# ---------------------------------------------------------------------
|
||
# Tools
|
||
# ---------------------------------------------------------------------
|
||
|
||
@self.mcp.tool()
|
||
def save(self, key: str, value: Any) -> bool:
|
||
"""Save a value under ``key`` in the default namespace.
|
||
|
||
Parameters
|
||
----------
|
||
key: str
|
||
Identifier for the value.
|
||
value: Any
|
||
JSON‑serialisable value.
|
||
|
||
Returns
|
||
-------
|
||
bool
|
||
``True`` if the value was stored.
|
||
"""
|
||
data = self._load()
|
||
full_key = f"default:{key}"
|
||
data[full_key] = {"value": value, "timestamp": datetime.utcnow().isoformat()}
|
||
self._save(data)
|
||
return True
|
||
|
||
@self.mcp.tool()
|
||
def get(self, key: str) -> Optional[Dict[str, Any]]:
|
||
"""Retrieve a value from the default namespace.
|
||
|
||
Parameters
|
||
----------
|
||
key: str
|
||
Identifier to look up.
|
||
|
||
Returns
|
||
-------
|
||
dict | None
|
||
Dictionary with ``key``, ``value`` and ``timestamp`` or ``None``.
|
||
"""
|
||
data = self._load()
|
||
full_key = f"default:{key}"
|
||
entry = data.get(full_key)
|
||
if entry is None:
|
||
return None
|
||
return {"key": full_key, "value": entry["value"], "timestamp": entry["timestamp"]}
|
||
|
||
@self.mcp.tool()
|
||
def delete(self, key: str) -> bool:
|
||
"""Delete a key from the default namespace.
|
||
|
||
Parameters
|
||
----------
|
||
key: str
|
||
Identifier to delete.
|
||
|
||
Returns
|
||
-------
|
||
bool
|
||
``True`` if the key existed and was removed.
|
||
"""
|
||
data = self._load()
|
||
full_key = f"default:{key}"
|
||
if full_key in data:
|
||
del data[full_key]
|
||
self._save(data)
|
||
return True
|
||
return False
|
||
|
||
@self.mcp.tool()
|
||
def list_keys(self, pattern: str = "*") -> List[str]:
|
||
"""Return all keys matching a wildcard pattern.
|
||
|
||
Parameters
|
||
----------
|
||
pattern: str, optional
|
||
Wildcard pattern (``*`` and ``?`` supported). Defaults to ``*``.
|
||
|
||
Returns
|
||
-------
|
||
list[str]
|
||
List of matching keys.
|
||
"""
|
||
data = self._load()
|
||
return [k for k in data.keys() if fnmatch.fnmatch(k, pattern)]
|
||
|
||
@self.mcp.tool()
|
||
def save_with_namespace(self, key: str, value: Any, namespace: str = "default") -> bool:
|
||
"""Save a value under a specific namespace.
|
||
|
||
Parameters
|
||
----------
|
||
key: str
|
||
Identifier.
|
||
value: Any
|
||
JSON‑serialisable value.
|
||
namespace: str, optional
|
||
Namespace prefix. Defaults to ``default``.
|
||
|
||
Returns
|
||
-------
|
||
bool
|
||
``True`` if stored.
|
||
"""
|
||
data = self._load()
|
||
full_key = f"{namespace}:{key}"
|
||
data[full_key] = {"value": value, "timestamp": datetime.utcnow().isoformat()}
|
||
self._save(data)
|
||
return True
|
||
|
||
@self.mcp.tool()
|
||
def get_by_namespace(self, namespace: str = "default") -> List[Dict[str, Any]]:
|
||
"""Return all entries belonging to a namespace.
|
||
|
||
Parameters
|
||
----------
|
||
namespace: str, optional
|
||
Namespace to query. Defaults to ``default``.
|
||
|
||
Returns
|
||
-------
|
||
list[dict]
|
||
List of dictionaries with ``key``, ``value`` and ``timestamp``.
|
||
"""
|
||
data = self._load()
|
||
prefix = f"{namespace}:"
|
||
result: List[Dict[str, Any]] = []
|
||
for k, v in data.items():
|
||
if k.startswith(prefix):
|
||
result.append({"key": k, "value": v["value"], "timestamp": v["timestamp"]})
|
||
return result
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Client demonstration
|
||
# ---------------------------------------------------------------------------
|
||
|
||
async def client_demo() -> None:
|
||
"""Simple demo that exercises the server tools via FastMCP Client."""
|
||
client = Client("python main.py") # launch the same script as a server
|
||
await client.connect()
|
||
try:
|
||
# Save a value in default namespace
|
||
await client.call_tool("save", {"key": "user_name", "value": "Alex"})
|
||
# Retrieve it
|
||
res = await client.call_tool("get", {"key": "user_name"})
|
||
print("get default:", res)
|
||
# Save with namespace
|
||
await client.call_tool("save_with_namespace", {"key": "session_token", "value": "abc123", "namespace": "agent_1"})
|
||
# List all keys
|
||
keys = await client.call_tool("list_keys", {"pattern": "*"})
|
||
print("all keys:", keys)
|
||
# Get by namespace
|
||
ns = await client.call_tool("get_by_namespace", {"namespace": "agent_1"})
|
||
print("agent_1 namespace:", ns)
|
||
finally:
|
||
await client.close()
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Entry points
|
||
# ---------------------------------------------------------------------------
|
||
|
||
if __name__ == "__main__":
|
||
# If the script is launched without arguments, run the server.
|
||
import sys
|
||
if len(sys.argv) == 1:
|
||
server = MemoryServer()
|
||
server.mcp.run(transport="stdio", show_banner=False, log_level="ERROR")
|
||
else:
|
||
# Otherwise, run the client demo.
|
||
asyncio.run(client_demo())
|