memory_server.py updated
This commit is contained in:
+9
-24
@@ -3,7 +3,6 @@ import os
|
||||
from datetime import datetime
|
||||
from fastmcp import MCP, Tool
|
||||
|
||||
# In-memory store backed by JSON file
|
||||
DATA_FILE = "memory_store.json"
|
||||
|
||||
class MemoryStore:
|
||||
@@ -21,7 +20,6 @@ class MemoryStore:
|
||||
def _key(self, key: str, namespace: str = None) -> str:
|
||||
return f"{namespace}:{key}" if namespace else key
|
||||
|
||||
# CRUD operations
|
||||
def save(self, key: str, value: str):
|
||||
self.store[self._key(key)] = {"value": value, "timestamp": datetime.utcnow().isoformat()}
|
||||
self._save()
|
||||
@@ -34,9 +32,8 @@ class MemoryStore:
|
||||
self._save()
|
||||
|
||||
def list_keys(self):
|
||||
return [k for k in self.store.keys()]
|
||||
return list(self.store.keys())
|
||||
|
||||
# Namespace helpers
|
||||
def save_with_namespace(self, namespace: str, key: str, value: str):
|
||||
self.store[self._key(key, namespace)] = {"value": value, "timestamp": datetime.utcnow().isoformat()}
|
||||
self._save()
|
||||
@@ -44,28 +41,22 @@ class MemoryStore:
|
||||
def get_by_namespace(self, namespace: str, key: str) -> str:
|
||||
return self.store.get(self._key(key, namespace), {}).get("value")
|
||||
|
||||
# Instantiate store
|
||||
store = MemoryStore()
|
||||
|
||||
# Define tools for MCP
|
||||
@Tool(name="save", description="Save a value under a key")
|
||||
def save_tool(args):
|
||||
key = args["key"]
|
||||
value = args["value"]
|
||||
store.save(key, value)
|
||||
return f"Saved {key}"
|
||||
store.save(args["key"], args["value"])
|
||||
return f"Saved {args['key']}"
|
||||
|
||||
@Tool(name="get", description="Get value by key")
|
||||
def get_tool(args):
|
||||
key = args["key"]
|
||||
val = store.get(key)
|
||||
val = store.get(args["key"])
|
||||
return val if val is not None else "Key not found"
|
||||
|
||||
@Tool(name="delete", description="Delete a key")
|
||||
def delete_tool(args):
|
||||
key = args["key"]
|
||||
store.delete(key)
|
||||
return f"Deleted {key}"
|
||||
store.delete(args["key"])
|
||||
return f"Deleted {args['key']}"
|
||||
|
||||
@Tool(name="list_keys", description="List all keys")
|
||||
def list_keys_tool(_):
|
||||
@@ -73,20 +64,14 @@ def list_keys_tool(_):
|
||||
|
||||
@Tool(name="save_with_namespace", description="Save value under key in namespace")
|
||||
def save_ns_tool(args):
|
||||
ns = args["namespace"]
|
||||
key = args["key"]
|
||||
val = args["value"]
|
||||
store.save_with_namespace(ns, key, val)
|
||||
return f"Saved {ns}:{key}"
|
||||
store.save_with_namespace(args["namespace"], args["key"], args["value"])
|
||||
return f"Saved {args['namespace']}:{args['key']}"
|
||||
|
||||
@Tool(name="get_by_namespace", description="Get value by key in namespace")
|
||||
def get_ns_tool(args):
|
||||
ns = args["namespace"]
|
||||
key = args["key"]
|
||||
val = store.get_by_namespace(ns, key)
|
||||
val = store.get_by_namespace(args["namespace"], args["key"])
|
||||
return val if val is not None else "Key not found"
|
||||
|
||||
# Create MCP instance and register tools
|
||||
mcp = MCP()
|
||||
mcp.register_tool(save_tool)
|
||||
mcp.register_tool(get_tool)
|
||||
|
||||
Reference in New Issue
Block a user