add main.py
This commit is contained in:
@@ -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())
|
||||
Reference in New Issue
Block a user