137 lines
4.2 KiB
Python
137 lines
4.2 KiB
Python
import os
|
|
import json
|
|
import datetime
|
|
from pathlib import Path
|
|
from dotenv import load_dotenv
|
|
from fastmcp import FastMCP, mcp
|
|
from pydantic import BaseModel, Field
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
# Configuration
|
|
DATA_FILE = Path(os.getenv("MEMORY_DATA_FILE", "memory.json"))
|
|
|
|
# Pydantic model for stored value
|
|
class StoredItem(BaseModel):
|
|
key: str = Field(..., description="The key for the item")
|
|
value: str = Field(..., description="The stored value")
|
|
timestamp: str = Field(..., description="ISO 8601 timestamp when stored")
|
|
|
|
# In-memory cache
|
|
_memory_cache: dict[str, StoredItem] = {}
|
|
|
|
# Load persisted data on startup
|
|
def load_data() -> None:
|
|
global _memory_cache
|
|
if DATA_FILE.exists():
|
|
try:
|
|
with DATA_FILE.open("r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
for key, item in data.items():
|
|
try:
|
|
_memory_cache[key] = StoredItem(**item)
|
|
except Exception as exc:
|
|
print(f"[WARN] Failed to parse item {key}: {exc}")
|
|
except Exception as exc:
|
|
print(f"[ERROR] Unable to read data file {DATA_FILE}: {exc}")
|
|
|
|
# Persist data to disk
|
|
def persist_data() -> None:
|
|
try:
|
|
data = {k: v.dict() for k, v in _memory_cache.items()}
|
|
with DATA_FILE.open("w", encoding="utf-8") as f:
|
|
json.dump(data, f, indent=2)
|
|
except Exception as exc:
|
|
print(f"[ERROR] Unable to write data file {DATA_FILE}: {exc}")
|
|
|
|
# Tool definitions
|
|
@mcp.tool("save")
|
|
def save(key: str, value: str) -> str:
|
|
"""Save a key-value pair with timestamp."""
|
|
try:
|
|
item = StoredItem(key=key, value=value, timestamp=datetime.datetime.utcnow().isoformat())
|
|
_memory_cache[key] = item
|
|
persist_data()
|
|
return f"Saved key '{key}'."
|
|
except Exception as exc:
|
|
return f"Error saving key '{key}': {exc}"
|
|
|
|
@mcp.tool("get")
|
|
def get(key: str) -> str:
|
|
"""Retrieve value for a key."""
|
|
try:
|
|
item = _memory_cache[key]
|
|
return f"{item.value} (stored at {item.timestamp})"
|
|
except KeyError:
|
|
return f"Key '{key}' not found."
|
|
except Exception as exc:
|
|
return f"Error retrieving key '{key}': {exc}"
|
|
|
|
@mcp.tool("delete")
|
|
def delete(key: str) -> str:
|
|
"""Delete a key."""
|
|
try:
|
|
del _memory_cache[key]
|
|
persist_data()
|
|
return f"Deleted key '{key}'."
|
|
except KeyError:
|
|
return f"Key '{key}' not found."
|
|
except Exception as exc:
|
|
return f"Error deleting key '{key}': {exc}"
|
|
|
|
@mcp.tool("list_keys")
|
|
def list_keys() -> str:
|
|
"""List all keys."""
|
|
try:
|
|
keys = ", ".join(_memory_cache.keys()) or "No keys"
|
|
return f"Keys: {keys}"
|
|
except Exception as exc:
|
|
return f"Error listing keys: {exc}"
|
|
|
|
# Namespace tools
|
|
@mcp.tool("save_with_namespace")
|
|
def save_with_namespace(namespace: str, key: str, value: str) -> str:
|
|
"""Save a key-value pair under a namespace."""
|
|
namespaced_key = f"{namespace}:{key}"
|
|
return save(namespaced_key, value)
|
|
|
|
@mcp.tool("get_by_namespace")
|
|
def get_by_namespace(namespace: str, key: str) -> str:
|
|
"""Retrieve a value from a namespace."""
|
|
namespaced_key = f"{namespace}:{key}"
|
|
return get(namespaced_key)
|
|
|
|
# Initialize server
|
|
mcpServer = FastMCP()
|
|
# Register tools automatically via decorators
|
|
|
|
# Load existing data
|
|
load_data()
|
|
|
|
# Demo client if invoked with --client
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
if len(sys.argv) > 1 and sys.argv[1] == "--client":
|
|
# Simple interactive client using stdio transport
|
|
print("Running client demo. Type 'exit' to quit.")
|
|
print("Available tools: save, get, delete, list_keys, save_with_namespace, get_by_namespace")
|
|
while True:
|
|
try:
|
|
user_input = input("> ")
|
|
except EOFError:
|
|
break
|
|
if user_input.strip().lower() == "exit":
|
|
break
|
|
# Send command via mcp.run() with input string
|
|
try:
|
|
response = mcp.run(user_input)
|
|
print(response)
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|
|
else:
|
|
# Run server with stdio transport
|
|
print("Starting FastMCP server with stdio transport...")
|
|
mcp.run()
|