valkey-glide 2.2.2__cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

glide_shared/commands/server_modules/ft_options/ft_constants.py

@@ -0,0 +1,84 @@
+# Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0
+
+
+class CommandNames:
+    """
+    Command name constants for vector search.
+    """
+
+    FT_CREATE = "FT.CREATE"
+    FT_DROPINDEX = "FT.DROPINDEX"
+    FT_LIST = "FT._LIST"
+    FT_SEARCH = "FT.SEARCH"
+    FT_INFO = "FT.INFO"
+    FT_ALIASADD = "FT.ALIASADD"
+    FT_ALIASDEL = "FT.ALIASDEL"
+    FT_ALIASUPDATE = "FT.ALIASUPDATE"
+    FT_EXPLAIN = "FT.EXPLAIN"
+    FT_EXPLAINCLI = "FT.EXPLAINCLI"
+    FT_AGGREGATE = "FT.AGGREGATE"
+    FT_PROFILE = "FT.PROFILE"
+    FT_ALIASLIST = "FT._ALIASLIST"
+
+
+class FtCreateKeywords:
+    """
+    Keywords used in the FT.CREATE command.
+    """
+
+    SCHEMA = "SCHEMA"
+    AS = "AS"
+    SORTABLE = "SORTABLE"
+    UNF = "UNF"
+    NO_INDEX = "NOINDEX"
+    ON = "ON"
+    PREFIX = "PREFIX"
+    SEPARATOR = "SEPARATOR"
+    CASESENSITIVE = "CASESENSITIVE"
+    DIM = "DIM"
+    DISTANCE_METRIC = "DISTANCE_METRIC"
+    TYPE = "TYPE"
+    INITIAL_CAP = "INITIAL_CAP"
+    M = "M"
+    EF_CONSTRUCTION = "EF_CONSTRUCTION"
+    EF_RUNTIME = "EF_RUNTIME"
+
+
+class FtSearchKeywords:
+    """
+    Keywords used in the FT.SEARCH command.
+    """
+
+    RETURN = "RETURN"
+    TIMEOUT = "TIMEOUT"
+    PARAMS = "PARAMS"
+    LIMIT = "LIMIT"
+    COUNT = "COUNT"
+    AS = "AS"
+
+
+class FtAggregateKeywords:
+    """
+    Keywords used in the FT.AGGREGATE command.
+    """
+
+    LIMIT = "LIMIT"
+    FILTER = "FILTER"
+    GROUPBY = "GROUPBY"
+    REDUCE = "REDUCE"
+    AS = "AS"
+    SORTBY = "SORTBY"
+    MAX = "MAX"
+    APPLY = "APPLY"
+    LOAD = "LOAD"
+    TIMEOUT = "TIMEOUT"
+    PARAMS = "PARAMS"
+
+
+class FtProfileKeywords:
+    """
+    Keywords used in the FT.PROFILE command.
+    """
+
+    QUERY = "QUERY"
+    LIMITED = "LIMITED"

glide_shared/commands/server_modules/ft_options/ft_create_options.py

@@ -0,0 +1,423 @@
+# Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import List, Optional
+
+from glide_shared.commands.server_modules.ft_options.ft_constants import (
+    FtCreateKeywords,
+)
+from glide_shared.constants import TEncodable
+
+
+class FieldType(Enum):
+    """
+    All possible values for the data type of field identifier for the SCHEMA option.
+    """
+
+    TEXT = "TEXT"
+    """
+    If the field contains any blob of data.
+    """
+    TAG = "TAG"
+    """
+    If the field contains a tag field.
+    """
+    NUMERIC = "NUMERIC"
+    """
+    If the field contains a number.
+    """
+    VECTOR = "VECTOR"
+    """
+    If the field is a vector field that supports vector search.
+    """
+
+
+class VectorAlgorithm(Enum):
+    """
+    Algorithm for vector type fields used for vector similarity search.
+    """
+
+    HNSW = "HNSW"
+    """
+    Hierarchical Navigable Small World algorithm.
+    """
+    FLAT = "FLAT"
+    """
+    Flat algorithm or the brute force algorithm.
+    """
+
+
+class DistanceMetricType(Enum):
+    """
+    Distance metrics to measure the degree of similarity between two vectors.
+
+    The above metrics calculate distance between two vectors, where the smaller the value is, the
+    closer the two vectors are in the vector space.
+    """
+
+    L2 = "L2"
+    """
+    Euclidean distance
+    """
+    IP = "IP"
+    """
+    Inner product
+    """
+    COSINE = "COSINE"
+    """
+    Cosine distance
+    """
+
+
+class VectorType(Enum):
+    """
+    Type type for the vector field type.
+    """
+
+    FLOAT32 = "FLOAT32"
+    """
+    FLOAT32 type of vector. The only supported type.
+    """
+
+
+class Field(ABC):
+    """
+    Abstract base class for a vector search field.
+    """
+
+    @abstractmethod
+    def __init__(
+        self,
+        name: TEncodable,
+        type: FieldType,
+        alias: Optional[TEncodable] = None,
+    ):
+        """
+        Initialize a new field instance.
+
+        Args:
+            name (TEncodable): The name of the field.
+            type (FieldType): The type of the field.
+            alias (Optional[TEncodable]): An alias for the field.
+        """
+        self.name = name
+        self.type = type
+        self.alias = alias
+
+    @abstractmethod
+    def to_args(self) -> List[TEncodable]:
+        """
+        Get the arguments representing the field.
+
+        Returns:
+            List[TEncodable]: A list of field arguments.
+        """
+        args = [self.name]
+        if self.alias:
+            args.extend([FtCreateKeywords.AS, self.alias])
+        args.append(self.type.value)
+        return args
+
+
+class TextField(Field):
+    """
+    Field contains any blob of data.
+    """
+
+    def __init__(self, name: TEncodable, alias: Optional[TEncodable] = None):
+        """
+        Initialize a new TextField instance.
+
+        Args:
+            name (TEncodable): The name of the text field.
+            alias (Optional[TEncodable]): An alias for the field.
+        """
+        super().__init__(name, FieldType.TEXT, alias)
+
+    def to_args(self) -> List[TEncodable]:
+        args = super().to_args()
+        return args
+
+
+class TagField(Field):
+    """
+    Tag fields are similar to full-text fields, but they interpret the text as a simple list of
+    tags delimited by a separator character.
+
+    For `HASH fields, separator default is a comma `,`. For `JSON` fields, there is no
+    default separator; you must declare one explicitly if needed.
+    """
+
+    def __init__(
+        self,
+        name: TEncodable,
+        alias: Optional[TEncodable] = None,
+        separator: Optional[TEncodable] = None,
+        case_sensitive: bool = False,
+    ):
+        """
+        Initialize a new TagField instance.
+
+        Args:
+            name (TEncodable): The name of the tag field.
+            alias (Optional[TEncodable]): An alias for the field.
+            separator (Optional[TEncodable]): Specify how text in the attribute is split into individual tags. Must be a
+                single character.
+            case_sensitive (bool): Preserve the original letter cases of tags. If set to False, characters are converted to
+                lowercase by default.
+        """
+        super().__init__(name, FieldType.TAG, alias)
+        self.separator = separator
+        self.case_sensitive = case_sensitive
+
+    def to_args(self) -> List[TEncodable]:
+        args = super().to_args()
+        if self.separator:
+            args.extend([FtCreateKeywords.SEPARATOR, self.separator])
+        if self.case_sensitive:
+            args.append(FtCreateKeywords.CASESENSITIVE)
+        return args
+
+
+class NumericField(Field):
+    """
+    Field contains a number.
+    """
+
+    def __init__(self, name: TEncodable, alias: Optional[TEncodable] = None):
+        """
+        Initialize a new NumericField instance.
+
+        Args:
+            name (TEncodable): The name of the numeric field.
+            alias (Optional[TEncodable]): An alias for the field.
+        """
+        super().__init__(name, FieldType.NUMERIC, alias)
+
+    def to_args(self) -> List[TEncodable]:
+        args = super().to_args()
+        return args
+
+
+class VectorFieldAttributes(ABC):
+    """
+    Abstract base class for defining vector field attributes to be used after the vector algorithm name.
+    """
+
+    @abstractmethod
+    def __init__(
+        self, dimensions: int, distance_metric: DistanceMetricType, type: VectorType
+    ):
+        """
+        Initialize a new vector field attributes instance.
+
+        Args:
+            dimensions (int): Number of dimensions in the vector. Equivalent to `DIM` on the module API.
+            distance_metric (DistanceMetricType): The distance metric used in vector type field. Can be one of
+                `[L2 | IP | COSINE]`. Equivalent to `DISTANCE_METRIC` on the module API.
+            type (VectorType): Vector type. The only supported type is `FLOAT32`. Equivalent to `TYPE` on the module API.
+        """
+        self.dimensions = dimensions
+        self.distance_metric = distance_metric
+        self.type = type
+
+    @abstractmethod
+    def to_args(self) -> List[TEncodable]:
+        """
+        Get the arguments to be used for the algorithm of the vector field.
+
+        Returns:
+            List[TEncodable]: A list of arguments.
+        """
+        args: List[TEncodable] = []
+        if self.dimensions:
+            args.extend([FtCreateKeywords.DIM, str(self.dimensions)])
+        if self.distance_metric:
+            args.extend([FtCreateKeywords.DISTANCE_METRIC, self.distance_metric.name])
+        if self.type:
+            args.extend([FtCreateKeywords.TYPE, self.type.name])
+        return args
+
+
+class VectorFieldAttributesFlat(VectorFieldAttributes):
+    """
+    Get the arguments to be used for the FLAT algorithm of the vector field.
+    """
+
+    def __init__(
+        self,
+        dimensions: int,
+        distance_metric: DistanceMetricType,
+        type: VectorType,
+        initial_cap: Optional[int] = None,
+    ):
+        """
+        Initialize a new flat vector field attributes instance.
+
+        Args:
+            dimensions (int): Number of dimensions in the vector. Equivalent to `DIM` on the module API.
+            distance_metric (DistanceMetricType): The distance metric used in vector type field. Can be one of
+                `[L2 | IP | COSINE]`. Equivalent to `DISTANCE_METRIC` on the module API.
+            type (VectorType): Vector type. The only supported type is `FLOAT32`. Equivalent to `TYPE` on the module API.
+            initial_cap (Optional[int]): Initial vector capacity in the index affecting memory allocation size of the index.
+                Defaults to `1024`. Equivalent to `INITIAL_CAP` on the module API.
+        """
+        super().__init__(dimensions, distance_metric, type)
+        self.initial_cap = initial_cap
+
+    def to_args(self) -> List[TEncodable]:
+        args = super().to_args()
+        if self.initial_cap:
+            args.extend([FtCreateKeywords.INITIAL_CAP, str(self.initial_cap)])
+        return args
+
+
+class VectorFieldAttributesHnsw(VectorFieldAttributes):
+    """
+    Get the arguments to be used for the HNSW algorithm of the vector field.
+    """
+
+    def __init__(
+        self,
+        dimensions: int,
+        distance_metric: DistanceMetricType,
+        type: VectorType,
+        initial_cap: Optional[int] = None,
+        number_of_edges: Optional[int] = None,
+        vectors_examined_on_construction: Optional[int] = None,
+        vectors_examined_on_runtime: Optional[int] = None,
+    ):
+        """
+        Initialize a new HNSW vector field attributes instance.
+
+        Args:
+            dimensions (int): Number of dimensions in the vector. Equivalent to `DIM` on the module API.
+            distance_metric (DistanceMetricType): The distance metric used in vector type field. Can be one of
+                `[L2 | IP | COSINE]`. Equivalent to `DISTANCE_METRIC` on the module API.
+            type (VectorType): Vector type. The only supported type is `FLOAT32`. Equivalent to `TYPE` on the module API.
+            initial_cap (Optional[int]): Initial vector capacity in the index affecting memory allocation size of the index.
+                Defaults to `1024`. Equivalent to `INITIAL_CAP` on the module API.
+            number_of_edges (Optional[int]): Number of maximum allowed outgoing edges for each node in the graph in each layer.
+                Default is `16`, maximum is `512`. Equivalent to `M` on the module API.
+            vectors_examined_on_construction (Optional[int]): Controls the number of vectors examined during index
+                construction. Default value is `200`, Maximum value is `4096`. Equivalent to `EF_CONSTRUCTION` on the
+                module API.
+            vectors_examined_on_runtime (Optional[int]): Controls the number of vectors examined during query operations.
+                Default value is `10`, Maximum value is `4096`. Equivalent to `EF_RUNTIME` on the module API.
+        """
+        super().__init__(dimensions, distance_metric, type)
+        self.initial_cap = initial_cap
+        self.number_of_edges = number_of_edges
+        self.vectors_examined_on_construction = vectors_examined_on_construction
+        self.vectors_examined_on_runtime = vectors_examined_on_runtime
+
+    def to_args(self) -> List[TEncodable]:
+        args = super().to_args()
+        if self.initial_cap:
+            args.extend([FtCreateKeywords.INITIAL_CAP, str(self.initial_cap)])
+        if self.number_of_edges:
+            args.extend([FtCreateKeywords.M, str(self.number_of_edges)])
+        if self.vectors_examined_on_construction:
+            args.extend(
+                [
+                    FtCreateKeywords.EF_CONSTRUCTION,
+                    str(self.vectors_examined_on_construction),
+                ]
+            )
+        if self.vectors_examined_on_runtime:
+            args.extend(
+                [FtCreateKeywords.EF_RUNTIME, str(self.vectors_examined_on_runtime)]
+            )
+        return args
+
+
+class VectorField(Field):
+    """
+    Class for defining vector field in a schema.
+    """
+
+    def __init__(
+        self,
+        name: TEncodable,
+        algorithm: VectorAlgorithm,
+        attributes: VectorFieldAttributes,
+        alias: Optional[TEncodable] = None,
+    ):
+        """
+        Initialize a new VectorField instance.
+
+        Args:
+            name (TEncodable): The name of the vector field.
+            algorithm (VectorAlgorithm): The vector indexing algorithm.
+            alias (Optional[TEncodable]): An alias for the field.
+            attributes (VectorFieldAttributes): Additional attributes to be passed with the vector field after the
+                algorithm name.
+        """
+        super().__init__(name, FieldType.VECTOR, alias)
+        self.algorithm = algorithm
+        self.attributes = attributes
+
+    def to_args(self) -> List[TEncodable]:
+        args = super().to_args()
+        args.append(self.algorithm.value)
+        if self.attributes:
+            attribute_list = self.attributes.to_args()
+            args.append(str(len(attribute_list)))
+            args.extend(attribute_list)
+        return args
+
+
+class DataType(Enum):
+    """
+    Type of the index dataset.
+    """
+
+    HASH = "HASH"
+    """
+    Data stored in hashes, so field identifiers are field names within the hashes.
+    """
+    JSON = "JSON"
+    """
+    Data stored as a JSON document, so field identifiers are JSON Path expressions.
+    """
+
+
+class FtCreateOptions:
+    """
+    This class represents the input options to be used in the FT.CREATE command.
+    All fields in this class are optional inputs for FT.CREATE.
+    """
+
+    def __init__(
+        self,
+        data_type: Optional[DataType] = None,
+        prefixes: Optional[List[TEncodable]] = None,
+    ):
+        """
+        Initialize the FT.CREATE optional fields.
+
+        Args:
+            data_type (Optional[DataType]): The index data type. If not defined a `HASH` index is created.
+            prefixes (Optional[List[TEncodable]]): A list of prefixes of index definitions.
+        """
+        self.data_type = data_type
+        self.prefixes = prefixes
+
+    def to_args(self) -> List[TEncodable]:
+        """
+        Get the optional arguments for the FT.CREATE command.
+
+        Returns:
+            List[TEncodable]:
+                List of FT.CREATE optional agruments.
+        """
+        args: List[TEncodable] = []
+        if self.data_type:
+            args.append(FtCreateKeywords.ON)
+            args.append(self.data_type.value)
+        if self.prefixes:
+            args.append(FtCreateKeywords.PREFIX)
+            args.append(str(len(self.prefixes)))
+            for prefix in self.prefixes:
+                args.append(prefix)
+        return args

glide_shared/commands/server_modules/ft_options/ft_profile_options.py

@@ -0,0 +1,113 @@
+# Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0
+from enum import Enum
+from typing import List, Optional, Union, cast
+
+from glide_shared.commands.server_modules.ft_options.ft_aggregate_options import (
+    FtAggregateOptions,
+)
+from glide_shared.commands.server_modules.ft_options.ft_constants import (
+    FtProfileKeywords,
+)
+from glide_shared.commands.server_modules.ft_options.ft_search_options import (
+    FtSearchOptions,
+)
+from glide_shared.constants import TEncodable
+
+
+class QueryType(Enum):
+    """
+    This class represents the query type being profiled.
+    """
+
+    AGGREGATE = "AGGREGATE"
+    """
+    If the query being profiled is for the FT.AGGREGATE command.
+    """
+    SEARCH = "SEARCH"
+    """
+    If the query being profiled is for the FT.SEARCH command.
+    """
+
+
+class FtProfileOptions:
+    """
+    This class represents the arguments/options for the FT.PROFILE command.
+    """
+
+    def __init__(
+        self,
+        query: TEncodable,
+        query_type: QueryType,
+        query_options: Optional[Union[FtSearchOptions, FtAggregateOptions]] = None,
+        limited: Optional[bool] = False,
+    ):
+        """
+        Initialize a new FtProfileOptions instance.
+
+        Args:
+            query (TEncodable): The query that is being profiled. This is the query argument from the
+                FT.AGGREGATE/FT.SEARCH command.
+            query_type (Optional[QueryType]): The type of query to be profiled.
+            query_options (Optional[Union[FtSearchOptions, FtAggregateOptions]]): The arguments/options for the
+                FT.AGGREGATE/FT.SEARCH command being profiled.
+            limited (Optional[bool]): To provide some brief version of the output, otherwise a full verbose output is provided.
+        """
+        self.query = query
+        self.query_type = query_type
+        self.query_options = query_options
+        self.limited = limited
+
+    @classmethod
+    def from_query_options(
+        cls,
+        query: TEncodable,
+        query_options: Union[FtSearchOptions, FtAggregateOptions],
+        limited: Optional[bool] = False,
+    ):
+        """
+        A class method to create FtProfileOptions with FT.SEARCH/FT.AGGREGATE options.
+
+        Args:
+            query (TEncodable): The query that is being profiled. This is the query argument from the
+                FT.AGGREGATE/FT.SEARCH command.
+            query_options (Optional[Union[FtSearchOptions, FtAggregateOptions]]): The arguments/options for the
+                FT.AGGREGATE/FT.SEARCH command being profiled.
+            limited (Optional[bool]): To provide some brief version of the output, otherwise a full verbose output is provided.
+        """
+        query_type: QueryType = QueryType.SEARCH
+        if isinstance(query_options, FtAggregateOptions):
+            query_type = QueryType.AGGREGATE
+        return cls(query, query_type, query_options, limited)
+
+    @classmethod
+    def from_query_type(
+        cls, query: TEncodable, query_type: QueryType, limited: Optional[bool] = False
+    ):
+        """
+        A class method to create FtProfileOptions with QueryType.
+
+        Args:
+            query (TEncodable): The query that is being profiled. This is the query argument from the
+                FT.AGGREGATE/FT.SEARCH command.
+            query_type (QueryType): The type of query to be profiled.
+            limited (Optional[bool]): To provide some brief version of the output, otherwise a full verbose output is provided.
+        """
+        return cls(query, query_type, None, limited)
+
+    def to_args(self) -> List[TEncodable]:
+        """
+        Get the remaining arguments for the FT.PROFILE command.
+
+        Returns:
+            List[TEncodable]: A list of remaining arguments for the FT.PROFILE command.
+        """
+        args: List[TEncodable] = [self.query_type.value]
+        if self.limited:
+            args.append(FtProfileKeywords.LIMITED)
+        args.extend([FtProfileKeywords.QUERY, self.query])
+        if self.query_options:
+            if isinstance(self.query_options, FtAggregateOptions):
+                args.extend(cast(FtAggregateOptions, self.query_options).to_args())
+            else:
+                args.extend(cast(FtSearchOptions, self.query_options).to_args())
+        return args