diff --git a/main.py b/main.py new file mode 100644 index 0000000..1c31896 --- /dev/null +++ b/main.py @@ -0,0 +1,234 @@ +#!/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())