Source code for coredis.modules.search

from __future__ import annotations

import dataclasses
import itertools
from abc import ABC, abstractmethod
from collections import OrderedDict
from datetime import timedelta

from deprecated.sphinx import versionadded

from ..commands._routing import FanoutStrategy, RandomStrategy
from ..commands._utils import normalized_milliseconds, normalized_seconds
from ..commands._validators import mutually_exclusive_parameters, mutually_inclusive_parameters
from ..commands._wrappers import ClusterCommandConfig
from ..commands.constants import CommandGroup, CommandName, NodeFlag
from ..commands.request import CommandRequest
from ..response._callbacks import (
    AnyStrCallback,
    ClusterEnsureConsistent,
    ClusterMergeSets,
    DictCallback,
    IntCallback,
    SetCallback,
    SimpleStringCallback,
)
from ..tokens import PrefixToken, PureToken
from ..typing import (
    AnyStr,
    CommandArgList,
    Key,
    KeyT,
    Literal,
    Mapping,
    Parameters,
    ResponsePrimitive,
    ResponseType,
    StringT,
    ValueT,
)
from .base import Module, ModuleGroup, module_command
from .response._callbacks.search import (
    AggregationResultCallback,
    HybridSearchCallback,
    SearchConfigCallback,
    SearchResultCallback,
    SpellCheckCallback,
)
from .response.types import HybridResult, 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/develop/ai/search-and-query"




[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.GEOSHAPE,
        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: StringT | None = None
    #: Whether to optimize for sorting.
    sortable: bool | None = None
    #: Whether to use the unnormalized form of the field for sorting.
    unf: bool | None = None
    #: Whether to disable stemming for this field.
    nostem: bool | None = None
    #: Skip indexing this field
    noindex: bool | None = None
    #: Phonetic algorithm to use for this field.
    phonetic: StringT | None = None
    #: Weight of this field in the document's ranking. The default is 1.0.
    weight: int | float | None = None
    #: Separator to use for splitting tags if the field is of
    #: type :attr:`~coredis.PureToken.TAG`.
    separator: StringT | None = 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: bool | None = 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: bool | None = None
    #: The algorithm to use for indexing if the field is of type
    #: :attr:`~coredis.PureToken.VECTOR`.
    #: For more details refer to the
    #: `Vector search concepts <https://redis.io/docs/develop/ai/search-and-query/vectors>`__
    #: section of the RediSearch documentation.
    algorithm: Literal["FLAT", "HSNW"] | None = None
    #: A dictionary of attributes to be used with the :paramref:`algorithm` specified.
    #: For more details refer to the
    #: `Vector index <https://redis.io/docs/develop/ai/search-and-query/vectors/#create-a-vector-index>`__
    #: section of the RediSearch documentation.
    attributes: dict[StringT, ValueT] | None = None

    @property
    def args(self) -> tuple[ValueT | Key, ...]:
        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/develop/ai/search-and-query/advanced-concepts/aggregations/#groupby-reducers>`__
    in the RediSearch documentation.
    """

    #: The name of the reducer function
    function: StringT
    #: The arguments to the reducer function
    parameters: Parameters[ValueT] | None = None
    #: The alias to assign to the result of the reducer function
    alias: StringT | None = 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/latest/develop/ai/search-and-query/advanced-concepts/aggregations/>`__
    in the RediSearch documentation.
    """

    #: The field to group by
    by: StringT | Parameters[StringT]
    #: The reducers to apply to each group
    reducers: Parameters[Reduce] | None = 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(PureToken.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/develop/ai/search-and-query/advanced-concepts/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/develop/ai/search-and-query/advanced-concepts/aggregations/#filter-expressions>`__
    in the RediSearch documentation.
    """

    #: The filter expression
    expression: StringT

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




@dataclasses.dataclass
class Combine(ABC):
    """
    Method definition for fusing text & vector results to be used with
    :paramref:`~coredis.modules.Search.hybrid.combine`
    to specify the ``COMBINE`` method for :meth:`~coredis.modules.Search.hybrid`
    """

    method: Literal["RRF", "LINEAR"] = dataclasses.field(init=False)
    window: int | None = None
    score_alias: StringT | None = None

    @property
    def args(self) -> CommandArgList:
        args = [PureToken.COMBINE, self.method]
        sub_args = self.sub_args
        if self.window is not None:
            sub_args.extend([PrefixToken.WINDOW, self.window])
        if self.score_alias:
            sub_args.extend([PrefixToken.YIELD_SCORE_AS, self.score_alias])
        return args + [len(sub_args), *sub_args]

    @property
    @abstractmethod
    def sub_args(self) -> CommandArgList: ...




[docs]
@dataclasses.dataclass
class RRFCombine(Combine):
    """
    RRF (Reciprocal Rank Fusion) method
    to be used in :meth:`~coredis.modules.Search.hybrid`
    with :paramref:`~coredis.modules.Search.hybrid.combine`
    """

    method: Literal["RRF"] = dataclasses.field(init=False, default="RRF")
    constant: int | None = None

    @property
    def sub_args(self) -> CommandArgList:
        args: CommandArgList = []
        if self.constant is not None:
            args.extend([PrefixToken.CONSTANT, self.constant])
        return args






[docs]
@dataclasses.dataclass
class LinearCombine(Combine):
    """
    Linear combination with ALPHA and BETA weights
    to be used in :meth:`~coredis.modules.Search.hybrid`
    with :paramref:`~coredis.modules.Search.hybrid.combine`
    """

    method: Literal["LINEAR"] = dataclasses.field(init=False, default="LINEAR")
    alpha: int | None = None
    beta: int | None = None

    @property
    def sub_args(self) -> CommandArgList:
        args: CommandArgList = []
        if self.alpha is not None:
            args.extend([PrefixToken.ALPHA, self.alpha])
        if self.beta is not None:
            args.extend([PrefixToken.BETA, self.beta])
        return args