Source code for coredis.modules.search

from __future__ import annotations

import dataclasses
import itertools
from datetime import timedelta

from deprecated.sphinx import versionadded

from ..commands._utils import normalized_milliseconds, normalized_seconds
from ..commands._wrappers import ClusterCommandConfig
from ..commands.constants import CommandGroup, CommandName, NodeFlag
from ..response._callbacks import (
    AnyStrCallback,
    ClusterEnsureConsistent,
    ClusterMergeSets,
    DictCallback,
    IntCallback,
    SetCallback,
    SimpleStringCallback,
)
from ..tokens import PrefixToken, PureToken
from ..typing import (
    AnyStr,
    CommandArgList,
    Dict,
    KeyT,
    List,
    Literal,
    Mapping,
    Optional,
    Parameters,
    ResponsePrimitive,
    ResponseType,
    Set,
    StringT,
    Tuple,
    Union,
    ValueT,
)
from .base import Module, ModuleGroup, module_command
from .response._callbacks.search import (
    AggregationResultCallback,
    SearchConfigCallback,
    SearchResultCallback,
    SpellCheckCallback,
    SpellCheckResult,
)
from .response.types import SearchAggregationResult, SearchResult


class RediSearch(Module[AnyStr]):
    NAME = "search"
    FULL_NAME = "RediSearch"
    DESCRIPTION = """RedisSearch is a Redis module that enables querying, secondary 
indexing, and full-text search for Redis. These features enable multi-field queries, 
aggregation, exact phrase matching, numeric filtering, geo filtering and vector 
similarity semantic search on top of text queries."""
    DOCUMENTATION_URL = "https://redis.io/docs/stack/search"




[docs]
@dataclasses.dataclass
class Field:
    """
    Field definition to be used in :meth:`~coredis.modules.Search.create` &
    :meth:`~coredis.modules.Search.alter`

    For more details refer to the documentation for
    `FT.CREATE <https://redis.io/commands/ft.create/>`__
    """

    #: Name of the field. For hashes, this is a field name within the hash.
    #:  For JSON, this is a JSON Path expression.
    name: StringT
    #: Type of field
    type: Literal[
        PureToken.TEXT,
        PureToken.TAG,
        PureToken.NUMERIC,
        PureToken.GEO,
        PureToken.VECTOR,
    ]

    #: Defines the alias associated to :paramref:`name`.
    #: For example, you can use this feature to alias a complex
    #: JSONPath expression with more memorable (and easier to type) name.
    alias: Optional[StringT] = None
    #: Whether to optimize for sorting.
    sortable: Optional[bool] = None
    #: Whether to use the unnormalized form of the field for sorting.
    unf: Optional[bool] = None
    #: Whether to disable stemming for this field.
    nostem: Optional[bool] = None
    #: Skip indexing this field
    noindex: Optional[bool] = None
    #: Phonetic algorithm to use for this field.
    phonetic: Optional[StringT] = None
    #: Weight of this field in the document's ranking. The default is 1.0.
    weight: Optional[Union[int, float]] = None
    #: Separator to use for splitting tags if the field is of
    #: type :attr:`~coredis.PureToken.TAG`.
    separator: Optional[StringT] = None
    #: For fields of type :attr:`~coredis.PureToken.TAG`,
    #: keeps the original letter cases of the tags. If not specified,
    #: the characters are converted to lowercase.
    casesensitive: Optional[bool] = None
    #: For fields of type :attr:`~coredis.PureToken.TAG` &
    #: :attr:`~coredis.PureToken.TEXT`, keeps a suffix trie with all
    #: terms which match the suffix. It is used to optimize contains ``(foo)``
    #: and suffix ``(*foo)`` queries. Otherwise, a brute-force search on the trie
    #: is performed. If suffix trie exists for some fields, these queries will
    #: be disabled for other fields
    withsuffixtrie: Optional[bool] = None
    #: The algorithm to use for indexing if the field is of type
    #: :attr:`~coredis.PureToken.VECTOR`.
    #: For more details refer to the
    #: `Vector similarity <https://redis.io/docs/stack/search/reference/vectors/>`__
    #: section of the RediSearch documentation.
    algorithm: Optional[Literal["FLAT", "HSNW"]] = None
    #: A dictionary of attributes to be used with the :paramref:`algorithm` specified.
    #: For more details refer to the
    #: `Creation attributes per algorithm <https://redis.io/docs/stack/search/reference/vectors/#creation-attributes-per-algorithm>`__
    #: section of the RediSearch documentation.
    attributes: Optional[Dict[StringT, ValueT]] = None

    @property
    def args(self) -> Tuple[ValueT, ...]:
        args: CommandArgList = [self.name]
        if self.alias:
            args += [PrefixToken.AS, self.alias]
        args += [self.type]
        if self.type == PureToken.VECTOR:
            assert self.algorithm

            args += [self.algorithm]
            if self.attributes:
                _attributes: List[ValueT] = list(
                    itertools.chain(*self.attributes.items())
                )
                args += [len(_attributes)]
                args += _attributes

        if self.sortable:
            args += [PureToken.SORTABLE]
        if self.unf:
            args += [PureToken.UNF]
        if self.nostem:
            args += [b"NOSTEM"]
        if self.noindex:
            args += [PureToken.NOINDEX]
        if self.phonetic:
            args += [b"PHONETIC", self.phonetic]
        if self.weight:
            args += [b"WEIGHT", self.weight]
        if self.separator:
            args += [PrefixToken.SEPARATOR, self.separator]
        if self.casesensitive:
            args += [b"CASESENSITIVE"]
        if self.withsuffixtrie:
            args += [PureToken.WITHSUFFIXTRIE]
        return tuple(args)






[docs]
@dataclasses.dataclass
class Reduce:
    """
    Reduce definition to be used with :paramref:`~coredis.modules.Search.aggregate.transformations`
    to define ``REDUCE`` steps in :meth:`~coredis.modules.Search.aggregate`

    For more details refer to `GroupBy Reducers <https://redis.io/docs/stack/search/reference/aggregations/#groupby-reducers>`__
    in the RediSearch documentation.
    """

    #: The name of the reducer function
    function: StringT
    #: The arguments to the reducer function
    parameters: Optional[Parameters[ValueT]] = None
    #: The alias to assign to the result of the reducer function
    alias: Optional[StringT] = None

    @property
    def args(self) -> CommandArgList:
        args: CommandArgList = [self.function]
        if self.parameters:
            args.extend(self.parameters)
        if self.alias:
            args.extend([PrefixToken.AS, self.alias])
        return args






[docs]
@dataclasses.dataclass
class Group:
    """
    Group definition to be used with :paramref:`~coredis.modules.Search.aggregate.transformations`
    to specify ``GROUPBY`` steps in :meth:`~coredis.modules.Search.aggregate`

    For more details refer to
    `Aggregations <https://redis.io/docs/stack/search/reference/aggregations>`__
    in the RediSearch documentation.
    """

    #: The field to group by
    by: Union[StringT, Parameters[StringT]]
    #: The reducers to apply to each group
    reducers: Optional[Parameters[Reduce]] = None

    @property
    def args(self) -> CommandArgList:
        args: CommandArgList = [PrefixToken.GROUPBY]
        if isinstance(self.by, (bytes, str)):
            args.extend([1, self.by])
        else:
            bies: List[StringT] = list(self.by)
            args.extend([len(bies), *bies])
        for reducer in self.reducers or []:
            args.append(PrefixToken.REDUCE)
            args.extend(reducer.args)

        return args






[docs]
@dataclasses.dataclass
class Apply:
    """
    Apply definition to be used with :paramref:`~coredis.modules.Search.aggregate.transformations`
    to specify ``APPLY`` steps in :meth:`~coredis.modules.Search.aggregate`

    For more details refer to
    `APPLY expressions <https://redis.io/docs/stack/search/reference/aggregations/#apply-expressions>`__
    in the RediSearch documentation.
    """

    #: The expression to apply
    function: StringT
    #: The alias to assign to the result of the expression
    alias: StringT

    @property
    def args(self) -> CommandArgList:
        return [PrefixToken.APPLY, self.function, PrefixToken.AS, self.alias]






[docs]
@dataclasses.dataclass
class Filter:
    """
    Filter definition to be used with :paramref:`~coredis.modules.Search.aggregate.transformations`
    to specify ``FILTER`` steps in :meth:`~coredis.modules.Search.aggregate`

    For more details refer to
    `FILTER expressions <https://redis.io/docs/stack/search/reference/aggregations/#filter-expressions>`__
    in the RediSearch documentation.
    """

    #: The filter expression
    expression: StringT

    @property
    def args(self) -> CommandArgList:
        return [PrefixToken.FILTER, self.expression]