memory_server.py updated

memory_server.py

@@ -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"]
+    store.save(args["key"], args["value"])
-    value = args["value"]
+    return f"Saved {args['key']}"
-    store.save(key, value)
-    return f"Saved {key}"
 
 @Tool(name="get", description="Get value by key")
 def get_tool(args):
-    key = args["key"]
+    val = store.get(args["key"])
-    val = store.get(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(args["key"])
-    store.delete(key)
+    return f"Deleted {args['key']}"
-    return f"Deleted {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"]
+    store.save_with_namespace(args["namespace"], args["key"], args["value"])
-    key = args["key"]
+    return f"Saved {args['namespace']}:{args['key']}"
-    val = args["value"]
-    store.save_with_namespace(ns, key, val)
-    return f"Saved {ns}:{key}"
 
 @Tool(name="get_by_namespace", description="Get value by key in namespace")
 def get_ns_tool(args):
-    ns = args["namespace"]
+    val = store.get_by_namespace(args["namespace"], args["key"])
-    key = args["key"]
-    val = store.get_by_namespace(ns, 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)