Adding more MCP tools for data structures

Author: mortensi

src/common/server.py

@@ -3,6 +3,6 @@
 # Initialize FastMCP server
 mcp = FastMCP(
     "Redis MCP Server",
-    dependencies=["redis", "pydantic"]
+    dependencies=["redis", "dotenv"]
 )
 

src/main.py

@@ -3,7 +3,15 @@
 import tools.server_management
 import tools.misc 
 import tools.redis_query_engine
-import tools.hash  
+import tools.hash
+import tools.list
+import tools.string
+import tools.json
+import tools.sorted_set
+import tools.set
+import tools.stream
+import tools.pub_sub
+
 
 class RedisMCPServer:
     def __init__(self):

src/tools/hash.py

@@ -2,22 +2,29 @@
 from redis.exceptions import RedisError
 from common.server import mcp
 
+
 @mcp.tool()
-async def hset(name: str, key: str, value: str) -> str:
-    """Set a field in a hash stored at key.
-    
+async def hset(name: str, key: str, value: str, expire_seconds: int = None) -> str:
+    """Set a field in a hash stored at key with an optional expiration time.
+
     Args:
         name: The Redis hash key.
         key: The field name inside the hash.
         value: The value to set.
-    
+        expire_seconds: Optional; time in seconds after which the key should expire.
+
     Returns:
         A success message or an error message.
     """
     try:
         r = RedisConnectionManager.get_connection()
         r.hset(name, key, value)
-        return f"Field '{key}' set successfully in hash '{name}'."
+
+        if expire_seconds is not None:
+            r.expire(name, expire_seconds)
+
+        return f"Field '{key}' set successfully in hash '{name}'." + (
+            f" Expires in {expire_seconds} seconds." if expire_seconds else "")
     except RedisError as e:
         return f"Error setting field '{key}' in hash '{name}': {str(e)}"
 

src/tools/json.py

@@ -0,0 +1,86 @@
+from common.connection import RedisConnectionManager
+from redis.exceptions import RedisError
+from common.server import mcp
+
+
+@mcp.tool()
+async def json_set(name: str, path: str, value: str, expire_seconds: int = None) -> str:
+    """Set a JSON value in Redis at a given path with an optional expiration time.
+
+    Args:
+        name: The Redis key where the JSON document is stored.
+        path: The JSON path where the value should be set.
+        value: The JSON value to store.
+        expire_seconds: Optional; time in seconds after which the key should expire.
+
+    Returns:
+        A success message or an error message.
+    """
+    try:
+        r = RedisConnectionManager.get_connection()
+        r.json().set(name, path, value)
+
+        if expire_seconds is not None:
+            r.expire(name, expire_seconds)
+
+        return f"JSON value set at path '{path}' in '{name}'." + (
+            f" Expires in {expire_seconds} seconds." if expire_seconds else "")
+    except RedisError as e:
+        return f"Error setting JSON value at path '{path}' in '{name}': {str(e)}"
+
+
+@mcp.tool()
+async def json_get(name: str, path: str = '$') -> str:
+    """Retrieve a JSON value from Redis at a given path.
+
+    Args:
+        name: The Redis key where the JSON document is stored.
+        path: The JSON path to retrieve (default: root '$').
+
+    Returns:
+        The retrieved JSON value or an error message.
+    """
+    try:
+        r = RedisConnectionManager.get_connection()
+        value = r.json().get(name, path)
+        return value if value is not None else f"No data found at path '{path}' in '{name}'."
+    except RedisError as e:
+        return f"Error retrieving JSON value at path '{path}' in '{name}': {str(e)}"
+
+
+@mcp.tool()
+async def json_del(name: str, path: str = '$') -> str:
+    """Delete a JSON value from Redis at a given path.
+
+    Args:
+        name: The Redis key where the JSON document is stored.
+        path: The JSON path to delete (default: root '$').
+
+    Returns:
+        A success message or an error message.
+    """
+    try:
+        r = RedisConnectionManager.get_connection()
+        deleted = r.json().delete(name, path)
+        return f"Deleted JSON value at path '{path}' in '{name}'." if deleted else f"No JSON value found at path '{path}' in '{name}'."
+    except RedisError as e:
+        return f"Error deleting JSON value at path '{path}' in '{name}': {str(e)}"
+
+
+@mcp.tool()
+async def json_expire(name: str, expire_seconds: int) -> str:
+    """Set an expiration time for a Redis JSON document.
+
+    Args:
+        name: The Redis key where the JSON document is stored.
+        expire_seconds: Time in seconds after which the key should expire.
+
+    Returns:
+        A success message or an error message.
+    """
+    try:
+        r = RedisConnectionManager.get_connection()
+        success = r.expire(name, expire_seconds)
+        return f"Expiration set to {expire_seconds} seconds for JSON document '{name}'." if success else f"JSON document '{name}' does not exist."
+    except RedisError as e:
+        return f"Error setting expiration for JSON document '{name}': {str(e)}"

src/tools/list.py

@@ -3,51 +3,32 @@
 from common.server import mcp
 
 @mcp.tool()
-async def lpush(name: str, value: str) -> str:
-    """Push a value onto the left of a Redis list.
-    
-    Args:
-        name: The Redis list key.
-        value: The value to push.
-    
-    Returns:
-        A success message or an error message.
-    """
+async def lpush(name: str, value: str, expire: int = None) -> str:
+    """Push a value onto the left of a Redis list and optionally set an expiration time."""
     try:
         r = RedisConnectionManager.get_connection()
         r.lpush(name, value)
+        if expire:
+            r.expire(name, expire)
         return f"Value '{value}' pushed to the left of list '{name}'."
     except RedisError as e:
         return f"Error pushing value to list '{name}': {str(e)}"
 
 @mcp.tool()
-async def rpush(name: str, value: str) -> str:
-    """Push a value onto the right of a Redis list.
-    
-    Args:
-        name: The Redis list key.
-        value: The value to push.
-    
-    Returns:
-        A success message or an error message.
-    """
+async def rpush(name: str, value: str, expire: int = None) -> str:
+    """Push a value onto the right of a Redis list and optionally set an expiration time."""
     try:
         r = RedisConnectionManager.get_connection()
         r.rpush(name, value)
+        if expire:
+            r.expire(name, expire)
         return f"Value '{value}' pushed to the right of list '{name}'."
     except RedisError as e:
         return f"Error pushing value to list '{name}': {str(e)}"
 
 @mcp.tool()
 async def lpop(name: str) -> str:
-    """Remove and return the first element from a Redis list.
-    
-    Args:
-        name: The Redis list key.
-    
-    Returns:
-        The removed element or an error message.
-    """
+    """Remove and return the first element from a Redis list."""
     try:
         r = RedisConnectionManager.get_connection()
         value = r.lpop(name)
@@ -57,14 +38,7 @@ async def lpop(name: str) -> str:
 
 @mcp.tool()
 async def rpop(name: str) -> str:
-    """Remove and return the last element from a Redis list.
-    
-    Args:
-        name: The Redis list key.
-    
-    Returns:
-        The removed element or an error message.
-    """
+    """Remove and return the last element from a Redis list."""
     try:
         r = RedisConnectionManager.get_connection()
         value = r.rpop(name)
@@ -74,16 +48,7 @@ async def rpop(name: str) -> str:
 
 @mcp.tool()
 async def lrange(name: str, start: int, stop: int) -> list:
-    """Get elements from a Redis list within a specific range.
-    
-    Args:
-        name: The Redis list key.
-        start: The starting index.
-        stop: The stopping index.
-    
-    Returns:
-        A list of elements or an error message.
-    """
+    """Get elements from a Redis list within a specific range."""
     try:
         r = RedisConnectionManager.get_connection()
         values = r.lrange(name, start, stop)
@@ -93,14 +58,7 @@ async def lrange(name: str, start: int, stop: int) -> list:
 
 @mcp.tool()
 async def llen(name: str) -> int:
-    """Get the length of a Redis list.
-    
-    Args:
-        name: The Redis list key.
-    
-    Returns:
-        The length of the list or an error message.
-    """
+    """Get the length of a Redis list."""
     try:
         r = RedisConnectionManager.get_connection()
         return r.llen(name)

src/tools/pub_sub.py

@@ -0,0 +1,60 @@
+from common.connection import RedisConnectionManager
+from redis.exceptions import RedisError
+from common.server import mcp
+
+
+@mcp.tool()
+async def publish(channel: str, message: str) -> str:
+    """Publish a message to a Redis channel.
+
+    Args:
+        channel: The Redis channel to publish to.
+        message: The message to send.
+
+    Returns:
+        A success message or an error message.
+    """
+    try:
+        r = RedisConnectionManager.get_connection()
+        r.publish(channel, message)
+        return f"Message published to channel '{channel}'."
+    except RedisError as e:
+        return f"Error publishing message to channel '{channel}': {str(e)}"
+
+
+@mcp.tool()
+async def subscribe(channel: str) -> str:
+    """Subscribe to a Redis channel.
+
+    Args:
+        channel: The Redis channel to subscribe to.
+
+    Returns:
+        A success message or an error message.
+    """
+    try:
+        r = RedisConnectionManager.get_connection()
+        pubsub = r.pubsub()
+        pubsub.subscribe(channel)
+        return f"Subscribed to channel '{channel}'."
+    except RedisError as e:
+        return f"Error subscribing to channel '{channel}': {str(e)}"
+
+
+@mcp.tool()
+async def unsubscribe(channel: str) -> str:
+    """Unsubscribe from a Redis channel.
+
+    Args:
+        channel: The Redis channel to unsubscribe from.
+
+    Returns:
+        A success message or an error message.
+    """
+    try:
+        r = RedisConnectionManager.get_connection()
+        pubsub = r.pubsub()
+        pubsub.unsubscribe(channel)
+        return f"Unsubscribed from channel '{channel}'."
+    except RedisError as e:
+        return f"Error unsubscribing from channel '{channel}': {str(e)}"

src/tools/redis_query_engine.py

@@ -2,6 +2,7 @@
 from common.connection import RedisConnectionManager
 from redis.exceptions import RedisError
 from common.server import mcp
+from redis.commands.search.query import Query
 
 
 @mcp.tool() 
@@ -12,4 +13,56 @@ async def get_indexes() -> str:
         r = RedisConnectionManager.get_connection()
         return r.execute_command("FT._LIST")
     except RedisError as e:
-        return f"Error retrieving indexes: {str(e)}"
+        return f"Error retrieving indexes: {str(e)}"
+
+
+@mcp.tool()
+async def get_index_info(index_name: str) -> str:
+    """Retrieve schema and information about a specific Redis index using FT.INFO.
+
+    Args:
+        index_name (str): The name of the index to retrieve information about.
+
+    Returns:
+        str: Information about the specified index or an error message.
+    """
+    try:
+        r = RedisConnectionManager.get_connection()
+        return r.ft(index_name).info()
+    except RedisError as e:
+        return f"Error retrieving index info: {str(e)}"
+
+
+@mcp.tool()
+async def get_indexed_keys(index_name: str) -> str:
+    """Retrieve the number of indexed keys by the index
+
+    Args:
+        index_name (str): The name of the index to retrieve information about.
+
+    Returns:
+        int: Number of indexed keys
+    """
+    try:
+        r = RedisConnectionManager.get_connection()
+        return r.ft(index_name).search(Query("*")).total
+    except RedisError as e:
+        return f"Error retrieving number of keys: {str(e)}"
+
+
+@mcp.tool()
+async def execute_redis_search(command: str, *args) -> str:
+    """Execute a raw Redis search command like FT.SEARCH or FT.AGGREGATE.
+
+    Args:
+        command (str): The Redis command to execute.
+        *args: Additional arguments for the command.
+
+    Returns:
+        str: The result of the command execution or an error message.
+    """
+    try:
+        r = RedisConnectionManager.get_connection()
+        return r.execute_command(command, *args)
+    except RedisError as e:
+        return f"Error executing command {command}: {str(e)}"