support for index creation and vector search

Author: mortensi

src/tools/misc.py

@@ -14,19 +14,6 @@ async def get_key_info(key: str) -> Dict[str, Any]:
             'type': key_type,
             'ttl': r.ttl(key)
         }
-
-        if key_type == 'string':
-            info['value'] = r.get(key)
-        elif key_type == 'list':
-            info['value'] = r.lrange(key, 0, -1)
-        elif key_type == 'hash':
-            info['value'] = r.hgetall(key)
-        elif key_type == 'set':
-            info['value'] = r.smembers(key)
-        elif key_type == 'zset':
-            info['value'] = r.zrange(key, 0, -1, withscores=True)
-        elif key_type == 'ReJSON-RL':
-            info['value'] = r.json().get(key, "$")
 
         return info
     except RedisError as e:

src/tools/redis_query_engine.py

@@ -3,6 +3,9 @@
 from redis.exceptions import RedisError
 from common.server import mcp
 from redis.commands.search.query import Query
+from redis.commands.search.field import VectorField
+from redis.commands.search.indexDefinition import IndexDefinition
+import numpy as np
 
 
 @mcp.tool() 
@@ -49,3 +52,84 @@ async def get_indexed_keys_number(index_name: str) -> str:
     except RedisError as e:
         return f"Error retrieving number of keys: {str(e)}"
 
+
+@mcp.tool()
+async def create_vector_index_hash(index_name: str,
+                       prefix: str = "doc:",
+                       vector_field: str = "vector",
+                       dim: int = 1536,
+                       distance_metric: str = "COSINE") -> str:
+    """
+    Create a Redis 8 vector similarity index using HNSW on a Redis hash.
+
+    This function sets up a Redis index for approximate nearest neighbor (ANN)
+    search using the HNSW algorithm and float32 vector embeddings.
+
+    Args:
+        index_name: The name of the Redis index to create.
+        prefix: The key prefix used to identify documents to index (e.g., 'doc:').
+        vector_field: The name of the vector field to be indexed for similarity search.
+        dim: The dimensionality of the vectors stored under the vector_field.
+        distance_metric: The distance function to use (e.g., 'COSINE', 'L2', 'IP').
+
+    Returns:
+        A string indicating whether the index was created successfully or an error message.
+    """
+    try:
+        r = RedisConnectionManager.get_connection()
+
+        index_def = IndexDefinition(prefix=[prefix])
+        schema = (
+            VectorField(
+                vector_field,
+                "HNSW",
+                {
+                    "TYPE": "FLOAT32",
+                    "DIM": dim,
+                    "DISTANCE_METRIC": distance_metric
+                }
+            )
+        )
+
+        r.ft(index_name).create_index([schema], definition=index_def)
+        return f"Index '{index_name}' created successfully."
+    except RedisError as e:
+        return f"Error creating index '{index_name}': {str(e)}"
+
+
+@mcp.tool()
+async def vector_search_hash(index_name: str,
+                            query_vector: list,
+                            vector_field: str = "vector",
+                            k: int = 5,
+                            return_fields: list = None) -> list:
+    """
+    Perform a KNN vector similarity search using Redis 8 or later version on vectors stored in hash data structures.
+
+    Args:
+        index_name: Name of the Redis index.
+        vector_field: Name of the indexed vector field.
+        query_vector: List of floats to use as the query vector.
+        k: Number of nearest neighbors to return.
+        return_fields: List of fields to return (optional).
+
+    Returns:
+        A list of matched documents or an error message.
+    """
+    try:
+        r = RedisConnectionManager.get_connection()
+
+        # Convert query vector to float32 binary blob
+        vector_blob = np.array(query_vector, dtype=np.float32).tobytes()
+
+        # Build the KNN query
+        base_query = f"*=>[KNN {k} @{vector_field} $vec_param AS score]"
+        query = Query(base_query).sort_by("score").paging(0, k).return_fields("id", "score", *return_fields or []).dialect(2)
+
+        # Perform the search with vector parameter
+        results = r.ft(index_name).search(query, query_params={"vec_param": vector_blob})
+
+        # Format and return the results
+        return [doc.__dict__ for doc in results.docs]
+    except RedisError as e:
+        return f"Error performing vector search on index '{index_name}': {str(e)}"