awslabs.valkey-mcp-server 0.1.1__py3-none-any.whl

awslabs/__init__.py

@@ -0,0 +1,13 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+
+# This file is part of the awslabs namespace.
+# It is intentionally minimal to support PEP 420 namespace packages.

awslabs/valkey_mcp_server/__init__.py

@@ -0,0 +1,14 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+
+"""Valkey MCP server package."""
+
+__version__ = '0.1.0'

awslabs/valkey_mcp_server/common/__init__.py

@@ -0,0 +1,20 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+
+"""
+Common imports.
+"""
+
+from . import (
+    config,
+    connection,
+    server,
+)

awslabs/valkey_mcp_server/common/config.py

@@ -0,0 +1,85 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+
+import os
+import urllib.parse
+from dotenv import load_dotenv
+
+
+load_dotenv()
+
+MCP_TRANSPORT = os.getenv('MCP_TRANSPORT', 'stdio')
+
+VALKEY_CFG = {
+    'host': os.getenv('VALKEY_HOST', '127.0.0.1'),
+    'port': int(os.getenv('VALKEY_PORT', 6379)),
+    'username': os.getenv('VALKEY_USERNAME', None),
+    'password': os.getenv('VALKEY_PWD', ''),
+    'ssl': os.getenv('VALKEY_USE_SSL', False) in ('true', '1', 't'),
+    'ssl_ca_path': os.getenv('VALKEY_SSL_CA_PATH', None),
+    'ssl_keyfile': os.getenv('VALKEY_SSL_KEYFILE', None),
+    'ssl_certfile': os.getenv('VALKEY_SSL_CERTFILE', None),
+    'ssl_cert_reqs': os.getenv('VALKEY_SSL_CERT_REQS', 'required'),
+    'ssl_ca_certs': os.getenv('VALKEY_SSL_CA_CERTS', None),
+    'cluster_mode': os.getenv('VALKEY_CLUSTER_MODE', False) in ('true', '1', 't'),
+}
+
+
+def generate_valkey_uri():
+    """Generates Valkey URL."""
+    cfg = VALKEY_CFG
+    scheme = 'valkeys' if cfg.get('ssl') else 'valkey'
+    host = cfg.get('host', '127.0.0.1')
+    port = cfg.get('port', 6379)
+
+    username = cfg.get('username')
+    password = cfg.get('password')
+
+    # Auth part - use quote() for auth components to preserve spaces as %20
+    def safe_quote(value):
+        """Safely quote a value that might be None."""
+        if value is None:
+            return ''
+        return urllib.parse.quote(str(value))
+
+    if username:
+        auth_part = f'{safe_quote(username)}:{safe_quote(password)}@'
+    elif password:
+        auth_part = f':{safe_quote(password)}@'
+    else:
+        auth_part = ''
+
+    # Base URI
+    base_uri = f'{scheme}://{auth_part}{host}:{port}'
+
+    # Additional SSL query parameters if SSL is enabled
+    query_params = {}
+    if cfg.get('ssl'):
+        if cfg.get('ssl_cert_reqs'):
+            query_params['ssl_cert_reqs'] = cfg['ssl_cert_reqs']
+        if cfg.get('ssl_ca_certs'):
+            query_params['ssl_ca_certs'] = cfg['ssl_ca_certs']
+        if cfg.get('ssl_keyfile'):
+            query_params['ssl_keyfile'] = cfg['ssl_keyfile']
+        if cfg.get('ssl_certfile'):
+            query_params['ssl_certfile'] = cfg['ssl_certfile']
+        if cfg.get('ssl_ca_path'):
+            query_params['ssl_ca_path'] = cfg['ssl_ca_path']
+
+    if query_params:
+        # Build query string with proper URL encoding
+        query_parts = []
+        for key, value in sorted(query_params.items()):
+            encoded_value = urllib.parse.quote(str(value), safe='')
+            query_parts.append(f'{key}={encoded_value}')
+        base_uri += '?' + '&'.join(query_parts)
+
+    return base_uri

awslabs/valkey_mcp_server/common/connection.py

@@ -0,0 +1,97 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+
+import sys
+from awslabs.valkey_mcp_server.common.config import VALKEY_CFG
+from awslabs.valkey_mcp_server.version import __version__
+from typing import Optional, Type, Union
+from valkey import (
+    Valkey,
+    exceptions,
+)
+from valkey.cluster import ValkeyCluster
+
+
+class ValkeyConnectionManager:
+    """Manages connection to Valkey."""
+
+    _instance: Optional[Union[Valkey, ValkeyCluster]] = None
+
+    @classmethod
+    def get_connection(cls, decode_responses: bool = True) -> Union[Valkey, ValkeyCluster]:
+        """Create connection to Valkey if none present or returns existing connection.
+
+        Args:
+            decode_responses: Whether to decode response bytes to strings. Defaults to True.
+
+        Returns:
+            Valkey: A Valkey connection instance.
+        """
+        if cls._instance is None:
+            try:
+                valkey_class: Type[Union[Valkey, ValkeyCluster]] = (
+                    ValkeyCluster if VALKEY_CFG['cluster_mode'] else Valkey
+                )
+
+                # Get SSL settings with defaults
+                ssl_enabled = VALKEY_CFG.get('ssl', False)
+                ssl_cert_reqs = VALKEY_CFG.get('ssl_cert_reqs')
+                if ssl_enabled and ssl_cert_reqs is None:
+                    ssl_cert_reqs = 'required'
+
+                # Build connection kwargs
+                connection_kwargs = {
+                    'host': VALKEY_CFG['host'],
+                    'port': VALKEY_CFG['port'],
+                    'username': VALKEY_CFG.get('username'),
+                    'password': VALKEY_CFG.get('password', ''),
+                    'ssl': ssl_enabled,
+                    'ssl_ca_path': VALKEY_CFG.get('ssl_ca_path'),
+                    'ssl_keyfile': VALKEY_CFG.get('ssl_keyfile'),
+                    'ssl_certfile': VALKEY_CFG.get('ssl_certfile'),
+                    'ssl_cert_reqs': ssl_cert_reqs,
+                    'ssl_ca_certs': VALKEY_CFG.get('ssl_ca_certs'),
+                    'decode_responses': decode_responses,
+                    'lib_name': f'valkey-py(mcp-server_v{__version__})',
+                }
+
+                # Add max_connections parameter based on mode
+                if VALKEY_CFG['cluster_mode']:
+                    connection_kwargs['max_connections_per_node'] = 10
+                else:
+                    connection_kwargs['max_connections'] = 10
+
+                # Create new instance
+                cls._instance = valkey_class(**connection_kwargs)
+
+            except exceptions.AuthenticationError:
+                print('Authentication failed', file=sys.stderr)
+                raise
+            except exceptions.ConnectionError:
+                print('Failed to connect to Valkey server', file=sys.stderr)
+                raise
+            except exceptions.TimeoutError:
+                print('Connection timed out', file=sys.stderr)
+                raise
+            except exceptions.ResponseError as e:
+                print(f'Response error: {e}', file=sys.stderr)
+                raise
+            except exceptions.ClusterError as e:
+                print(f'Valkey Cluster error: {e}', file=sys.stderr)
+                raise
+            except exceptions.ValkeyError as e:
+                print(f'Valkey error: {e}', file=sys.stderr)
+                raise
+            except Exception as e:
+                print(f'Unexpected error: {e}', file=sys.stderr)
+                raise
+
+        return cls._instance

awslabs/valkey_mcp_server/common/server.py

@@ -0,0 +1,23 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+
+from mcp.server.fastmcp import FastMCP
+
+
+# Initialize FastMCP server
+mcp = FastMCP(
+    'awslabs.valkey-mcp-server',
+    instructions='Instructions for using this valkey MCP server. This can be used by clients to improve the LLM'
+    's understanding of available tools, resources, etc. It can be thought of like a '
+    'hint'
+    ' to the model. For example, this information MAY be added to the system prompt. Important to be clear, direct, and detailed.',
+    dependencies=['pydantic', 'loguru', 'valkey', 'dotenv', 'numpy'],
+)

awslabs/valkey_mcp_server/main.py

@@ -0,0 +1,84 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+
+"""awslabs valkey MCP Server implementation."""
+
+import argparse
+from awslabs.valkey_mcp_server.common.server import mcp
+from awslabs.valkey_mcp_server.tools import (
+    bitmap,  # noqa: F401
+    hash,  # noqa: F401
+    hyperloglog,  # noqa: F401
+    json,  # noqa: F401
+    list,  # noqa: F401
+    misc,  # noqa: F401
+    server_management,  # noqa: F401
+    set,  # noqa: F401
+    sorted_set,  # noqa: F401
+    stream,  # noqa: F401
+    string,  # noqa: F401
+)
+from loguru import logger
+from starlette.requests import Request  # noqa: F401
+from starlette.responses import Response
+
+
+# Add a health check route directly to the MCP server
+@mcp.custom_route('/health', methods=['GET'])
+async def health_check(request):
+    """Simple health check endpoint for ALB Target Group.
+
+    Always returns 200 OK to indicate the service is running.
+    """
+    return Response(content='healthy', status_code=200, media_type='text/plain')
+
+
+class ValkeyMCPServer:
+    """Valkey MCP Server wrapper."""
+
+    def __init__(self, sse=False, port=8888):
+        """Initialize MCP Server wrapper.
+
+        Args:
+            sse: Whether to use SSE transport
+            port: Port to run the server on (default: 8888)
+        """
+        self.sse = sse
+        self.port = port
+
+    def run(self):
+        """Run server with appropriate transport."""
+        if self.sse:
+            if self.port is not None:
+                mcp.settings.port = self.port
+            mcp.run(transport='sse')
+        else:
+            mcp.run()
+
+
+def main():
+    """Run the MCP server with CLI argument support."""
+    parser = argparse.ArgumentParser(
+        description='An AWS Labs Model Context Protocol (MCP) server for valkey'
+    )
+    parser.add_argument('--sse', action='store_true', help='Use SSE transport')
+    parser.add_argument('--port', type=int, default=8888, help='Port to run the server on')
+
+    args = parser.parse_args()
+
+    logger.info('Amazon ElastiCache/MemoryDB Valkey MCP Server Started...')
+
+    server = ValkeyMCPServer(args.sse, args.port)
+    server.run()
+
+
+if __name__ == '__main__':
+    main()

awslabs/valkey_mcp_server/tools/__init__.py

@@ -0,0 +1,28 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+
+"""
+Tool imports for Valkey MCP Server.
+"""
+
+from . import (
+    bitmap,
+    hash,
+    hyperloglog,
+    json,
+    list,
+    misc,
+    server_management,
+    set,
+    sorted_set,
+    stream,
+    string,
+)

awslabs/valkey_mcp_server/tools/bitmap.py

@@ -0,0 +1,148 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+
+"""Bitmap operations for Valkey MCP Server."""
+
+from awslabs.valkey_mcp_server.common.connection import ValkeyConnectionManager
+from awslabs.valkey_mcp_server.common.server import mcp
+from typing import Optional
+from valkey.exceptions import ValkeyError
+
+
+@mcp.tool()
+async def bitmap_set(key: str, offset: int, value: int) -> str:
+    """Set the bit at offset to value.
+
+    Args:
+        key: The name of the bitmap key
+        offset: The bit offset (0-based)
+        value: The bit value (0 or 1)
+
+    Returns:
+        Success message or error message
+    """
+    try:
+        if value not in (0, 1):
+            return f'Error: value must be 0 or 1, got {value}'
+        if offset < 0:
+            return f'Error: offset must be non-negative, got {offset}'
+
+        r = ValkeyConnectionManager.get_connection()
+        previous = r.setbit(key, offset, value)
+        return f'Bit at offset {offset} set to {value} (previous value: {previous})'
+    except ValkeyError as e:
+        return f"Error setting bit in '{key}': {str(e)}"
+
+
+@mcp.tool()
+async def bitmap_get(key: str, offset: int) -> str:
+    """Get the bit value at offset.
+
+    Args:
+        key: The name of the bitmap key
+        offset: The bit offset (0-based)
+
+    Returns:
+        Bit value or error message
+    """
+    try:
+        if offset < 0:
+            return f'Error: offset must be non-negative, got {offset}'
+
+        r = ValkeyConnectionManager.get_connection()
+        value = r.getbit(key, offset)
+        return f'Bit at offset {offset} is {value}'
+    except ValkeyError as e:
+        return f"Error getting bit from '{key}': {str(e)}"
+
+
+@mcp.tool()
+async def bitmap_count(key: str, start: Optional[int] = None, end: Optional[int] = None) -> str:
+    """Count the number of set bits (1) in a range.
+
+    Args:
+        key: The name of the bitmap key
+        start: Start offset (inclusive, optional)
+        end: End offset (inclusive, optional)
+
+    Returns:
+        Count of set bits or error message
+    """
+    try:
+        r = ValkeyConnectionManager.get_connection()
+        if start is not None and end is not None:
+            if start < 0 or end < 0:
+                return 'Error: start and end must be non-negative'
+            if start > end:
+                return 'Error: start must be less than or equal to end'
+            count = r.bitcount(key, start, end)
+            range_str = f' in range [{start}, {end}]'
+        else:
+            count = r.bitcount(key)
+            range_str = ''
+
+        return f'Number of set bits{range_str}: {count}'
+    except ValkeyError as e:
+        return f"Error counting bits in '{key}': {str(e)}"
+
+
+@mcp.tool()
+async def bitmap_pos(
+    key: str,
+    bit: int,
+    start: Optional[int] = None,
+    end: Optional[int] = None,
+    count: Optional[int] = None,
+) -> str:
+    """Find positions of bits set to a specific value.
+
+    Args:
+        key: The name of the bitmap key
+        bit: Bit value to search for (0 or 1)
+        start: Start offset (inclusive, optional)
+        end: End offset (inclusive, optional)
+        count: Maximum number of positions to return (optional)
+
+    Returns:
+        List of positions or error message
+    """
+    try:
+        if bit not in (0, 1):
+            return f'Error: bit must be 0 or 1, got {bit}'
+
+        r = ValkeyConnectionManager.get_connection()
+        args = []
+        if start is not None:
+            if start < 0:
+                return 'Error: start must be non-negative'
+            args.extend(['START', start])
+        if end is not None:
+            if end < 0:
+                return 'Error: end must be non-negative'
+            if start is not None and start > end:
+                return 'Error: start must be less than or equal to end'
+            args.extend(['END', end])
+        if count is not None:
+            if count < 1:
+                return 'Error: count must be positive'
+            args.extend(['COUNT', count])
+
+        pos = r.bitpos(key, bit, *args) if args else r.bitpos(key, bit)
+
+        if pos == -1 or pos is None:
+            range_str = ''
+            if start is not None or end is not None:
+                range_str = f' in range [{start or 0}, {end or "∞"}]'
+            return f'No bits set to {bit} found{range_str}'
+
+        return f'First bit set to {bit} found at position: {pos}'
+    except ValkeyError as e:
+        return f"Error finding bit position in '{key}': {str(e)}"