elasticsearch 8.17.1__py3-none-any.whl → 9.0.0__py3-none-any.whl

elasticsearch/dsl/search_base.py

@@ -0,0 +1,1053 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+import collections.abc
+import copy
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    Iterator,
+    List,
+    Optional,
+    Protocol,
+    Tuple,
+    Type,
+    Union,
+    cast,
+    overload,
+)
+
+from typing_extensions import Self, TypeVar
+
+from .aggs import A, Agg, AggBase
+from .document_base import InstrumentedField
+from .exceptions import IllegalOperation
+from .query import Bool, Q, Query
+from .response import Hit, Response
+from .utils import _R, AnyUsingType, AttrDict, DslBase, recursive_to_dict
+
+if TYPE_CHECKING:
+    from .field import Field, Object
+
+
+class SupportsClone(Protocol):
+    def _clone(self) -> Self: ...
+
+
+_S = TypeVar("_S", bound=SupportsClone)
+
+
+class QueryProxy(Generic[_S]):
+    """
+    Simple proxy around DSL objects (queries) that can be called
+    (to add query/post_filter) and also allows attribute access which is proxied to
+    the wrapped query.
+    """
+
+    def __init__(self, search: _S, attr_name: str):
+        self._search = search
+        self._proxied: Optional[Query] = None
+        self._attr_name = attr_name
+
+    def __nonzero__(self) -> bool:
+        return self._proxied is not None
+
+    __bool__ = __nonzero__
+
+    def __call__(self, *args: Any, **kwargs: Any) -> _S:
+        """
+        Add a query.
+        """
+        s = self._search._clone()
+
+        # we cannot use self._proxied since we just cloned self._search and
+        # need to access the new self on the clone
+        proxied = getattr(s, self._attr_name)
+        if proxied._proxied is None:
+            proxied._proxied = Q(*args, **kwargs)
+        else:
+            proxied._proxied &= Q(*args, **kwargs)
+
+        # always return search to be chainable
+        return s
+
+    def __getattr__(self, attr_name: str) -> Any:
+        return getattr(self._proxied, attr_name)
+
+    def __setattr__(self, attr_name: str, value: Any) -> None:
+        if not attr_name.startswith("_"):
+            if self._proxied is not None:
+                self._proxied = Q(self._proxied.to_dict())
+                setattr(self._proxied, attr_name, value)
+        super().__setattr__(attr_name, value)
+
+    def __getstate__(self) -> Tuple[_S, Optional[Query], str]:
+        return self._search, self._proxied, self._attr_name
+
+    def __setstate__(self, state: Tuple[_S, Optional[Query], str]) -> None:
+        self._search, self._proxied, self._attr_name = state
+
+
+class ProxyDescriptor(Generic[_S]):
+    """
+    Simple descriptor to enable setting of queries and filters as:
+
+        s = Search()
+        s.query = Q(...)
+
+    """
+
+    def __init__(self, name: str):
+        self._attr_name = f"_{name}_proxy"
+
+    def __get__(self, instance: Any, owner: object) -> QueryProxy[_S]:
+        return cast(QueryProxy[_S], getattr(instance, self._attr_name))
+
+    def __set__(self, instance: _S, value: Dict[str, Any]) -> None:
+        proxy: QueryProxy[_S] = getattr(instance, self._attr_name)
+        proxy._proxied = Q(value)
+
+
+class AggsProxy(AggBase[_R], DslBase):
+    name = "aggs"
+
+    def __init__(self, search: "SearchBase[_R]"):
+        self._base = cast("Agg[_R]", self)
+        self._search = search
+        self._params = {"aggs": {}}
+
+    def to_dict(self) -> Dict[str, Any]:
+        return cast(Dict[str, Any], super().to_dict().get("aggs", {}))
+
+
+class Request(Generic[_R]):
+    def __init__(
+        self,
+        using: AnyUsingType = "default",
+        index: Optional[Union[str, List[str]]] = None,
+        doc_type: Optional[
+            Union[type, str, List[Union[type, str]], Dict[str, Union[type, str]]]
+        ] = None,
+        extra: Optional[Dict[str, Any]] = None,
+    ):
+        self._using = using
+
+        self._index = None
+        if isinstance(index, (tuple, list)):
+            self._index = list(index)
+        elif index:
+            self._index = [index]
+
+        self._doc_type: List[Union[type, str]] = []
+        self._doc_type_map: Dict[str, Any] = {}
+        if isinstance(doc_type, (tuple, list)):
+            self._doc_type.extend(doc_type)
+        elif isinstance(doc_type, collections.abc.Mapping):
+            self._doc_type.extend(doc_type.keys())
+            self._doc_type_map.update(doc_type)
+        elif doc_type:
+            self._doc_type.append(doc_type)
+
+        self._params: Dict[str, Any] = {}
+        self._extra: Dict[str, Any] = extra or {}
+
+    def __eq__(self, other: Any) -> bool:
+        return (
+            isinstance(other, Request)
+            and other._params == self._params
+            and other._index == self._index
+            and other._doc_type == self._doc_type
+            and other.to_dict() == self.to_dict()
+        )
+
+    def __copy__(self) -> Self:
+        return self._clone()
+
+    def params(self, **kwargs: Any) -> Self:
+        """
+        Specify query params to be used when executing the search. All the
+        keyword arguments will override the current values. See
+        https://elasticsearch-py.readthedocs.io/en/latest/api/elasticsearch.html#elasticsearch.Elasticsearch.search
+        for all available parameters.
+
+        Example::
+
+            s = Search()
+            s = s.params(routing='user-1', preference='local')
+        """
+        s = self._clone()
+        s._params.update(kwargs)
+        return s
+
+    def index(self, *index: Union[str, List[str], Tuple[str, ...]]) -> Self:
+        """
+        Set the index for the search. If called empty it will remove all information.
+
+        Example::
+
+            s = Search()
+            s = s.index('twitter-2015.01.01', 'twitter-2015.01.02')
+            s = s.index(['twitter-2015.01.01', 'twitter-2015.01.02'])
+        """
+        # .index() resets
+        s = self._clone()
+        if not index:
+            s._index = None
+        else:
+            indexes = []
+            for i in index:
+                if isinstance(i, str):
+                    indexes.append(i)
+                elif isinstance(i, list):
+                    indexes += i
+                elif isinstance(i, tuple):
+                    indexes += list(i)
+
+            s._index = (self._index or []) + indexes
+
+        return s
+
+    def _resolve_field(self, path: str) -> Optional["Field"]:
+        for dt in self._doc_type:
+            if not hasattr(dt, "_index"):
+                continue
+            field = dt._index.resolve_field(path)
+            if field is not None:
+                return cast("Field", field)
+        return None
+
+    def _resolve_nested(
+        self, hit: AttrDict[Any], parent_class: Optional[type] = None
+    ) -> Type[_R]:
+        doc_class = Hit
+
+        nested_path = []
+        nesting = hit["_nested"]
+        while nesting and "field" in nesting:
+            nested_path.append(nesting["field"])
+            nesting = nesting.get("_nested")
+        nested_path_str = ".".join(nested_path)
+
+        nested_field: Optional["Object"]
+        if parent_class is not None and hasattr(parent_class, "_index"):
+            nested_field = cast(
+                Optional["Object"], parent_class._index.resolve_field(nested_path_str)
+            )
+        else:
+            nested_field = cast(
+                Optional["Object"], self._resolve_field(nested_path_str)
+            )
+
+        if nested_field is not None:
+            return cast(Type[_R], nested_field._doc_class)
+
+        return cast(Type[_R], doc_class)
+
+    def _get_result(
+        self, hit: AttrDict[Any], parent_class: Optional[type] = None
+    ) -> _R:
+        doc_class: Any = Hit
+        dt = hit.get("_type")
+
+        if "_nested" in hit:
+            doc_class = self._resolve_nested(hit, parent_class)
+
+        elif dt in self._doc_type_map:
+            doc_class = self._doc_type_map[dt]
+
+        else:
+            for doc_type in self._doc_type:
+                if hasattr(doc_type, "_matches") and doc_type._matches(hit):
+                    doc_class = doc_type
+                    break
+
+        for t in hit.get("inner_hits", ()):
+            hit["inner_hits"][t] = Response[_R](
+                self, hit["inner_hits"][t], doc_class=doc_class
+            )
+
+        callback = getattr(doc_class, "from_es", doc_class)
+        return cast(_R, callback(hit))
+
+    def doc_type(
+        self, *doc_type: Union[type, str], **kwargs: Callable[[AttrDict[Any]], Any]
+    ) -> Self:
+        """
+        Set the type to search through. You can supply a single value or
+        multiple. Values can be strings or subclasses of ``Document``.
+
+        You can also pass in any keyword arguments, mapping a doc_type to a
+        callback that should be used instead of the Hit class.
+
+        If no doc_type is supplied any information stored on the instance will
+        be erased.
+
+        Example:
+
+            s = Search().doc_type('product', 'store', User, custom=my_callback)
+        """
+        # .doc_type() resets
+        s = self._clone()
+        if not doc_type and not kwargs:
+            s._doc_type = []
+            s._doc_type_map = {}
+        else:
+            s._doc_type.extend(doc_type)
+            s._doc_type.extend(kwargs.keys())
+            s._doc_type_map.update(kwargs)
+        return s
+
+    def using(self, client: AnyUsingType) -> Self:
+        """
+        Associate the search request with an elasticsearch client. A fresh copy
+        will be returned with current instance remaining unchanged.
+
+        :arg client: an instance of ``elasticsearch.Elasticsearch`` to use or
+            an alias to look up in ``elasticsearch.dsl.connections``
+
+        """
+        s = self._clone()
+        s._using = client
+        return s
+
+    def extra(self, **kwargs: Any) -> Self:
+        """
+        Add extra keys to the request body. Mostly here for backwards
+        compatibility.
+        """
+        s = self._clone()
+        if "from_" in kwargs:
+            kwargs["from"] = kwargs.pop("from_")
+        s._extra.update(kwargs)
+        return s
+
+    def _clone(self) -> Self:
+        s = self.__class__(
+            using=self._using, index=self._index, doc_type=self._doc_type
+        )
+        s._doc_type_map = self._doc_type_map.copy()
+        s._extra = self._extra.copy()
+        s._params = self._params.copy()
+        return s
+
+    if TYPE_CHECKING:
+
+        def to_dict(self) -> Dict[str, Any]: ...
+
+
+class SearchBase(Request[_R]):
+    query = ProxyDescriptor[Self]("query")
+    post_filter = ProxyDescriptor[Self]("post_filter")
+    _response: Response[_R]
+
+    def __init__(
+        self,
+        using: AnyUsingType = "default",
+        index: Optional[Union[str, List[str]]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Search request to elasticsearch.
+
+        :arg using: `Elasticsearch` instance to use
+        :arg index: limit the search to index
+
+        All the parameters supplied (or omitted) at creation type can be later
+        overridden by methods (`using`, `index` and `doc_type` respectively).
+        """
+        super().__init__(using=using, index=index, **kwargs)
+
+        self.aggs = AggsProxy[_R](self)
+        self._sort: List[Union[str, Dict[str, Dict[str, str]]]] = []
+        self._knn: List[Dict[str, Any]] = []
+        self._rank: Dict[str, Any] = {}
+        self._collapse: Dict[str, Any] = {}
+        self._source: Optional[Union[bool, List[str], Dict[str, List[str]]]] = None
+        self._highlight: Dict[str, Any] = {}
+        self._highlight_opts: Dict[str, Any] = {}
+        self._suggest: Dict[str, Any] = {}
+        self._script_fields: Dict[str, Any] = {}
+        self._response_class = Response[_R]
+
+        self._query_proxy = QueryProxy(self, "query")
+        self._post_filter_proxy = QueryProxy(self, "post_filter")
+
+    def filter(self, *args: Any, **kwargs: Any) -> Self:
+        """
+        Add a query in filter context.
+        """
+        return self.query(Bool(filter=[Q(*args, **kwargs)]))
+
+    def exclude(self, *args: Any, **kwargs: Any) -> Self:
+        """
+        Add a negative query in filter context.
+        """
+        return self.query(Bool(filter=[~Q(*args, **kwargs)]))
+
+    def __getitem__(self, n: Union[int, slice]) -> Self:
+        """
+        Support slicing the `Search` instance for pagination.
+
+        Slicing equates to the from/size parameters. E.g.::
+
+            s = Search().query(...)[0:25]
+
+        is equivalent to::
+
+            s = Search().query(...).extra(from_=0, size=25)
+
+        """
+        s = self._clone()
+
+        if isinstance(n, slice):
+            # If negative slicing, abort.
+            if n.start and n.start < 0 or n.stop and n.stop < 0:
+                raise ValueError("Search does not support negative slicing.")
+            slice_start = n.start
+            slice_stop = n.stop
+        else:  # This is an index lookup, equivalent to slicing by [n:n+1].
+            # If negative index, abort.
+            if n < 0:
+                raise ValueError("Search does not support negative indexing.")
+            slice_start = n
+            slice_stop = n + 1
+
+        old_from = s._extra.get("from")
+        old_to = None
+        if "size" in s._extra:
+            old_to = (old_from or 0) + s._extra["size"]
+
+        new_from = old_from
+        if slice_start is not None:
+            new_from = (old_from or 0) + slice_start
+        new_to = old_to
+        if slice_stop is not None:
+            new_to = (old_from or 0) + slice_stop
+            if old_to is not None and old_to < new_to:
+                new_to = old_to
+
+        if new_from is not None:
+            s._extra["from"] = new_from
+        if new_to is not None:
+            s._extra["size"] = max(0, new_to - (new_from or 0))
+        return s
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> Self:
+        """
+        Construct a new `Search` instance from a raw dict containing the search
+        body. Useful when migrating from raw dictionaries.
+
+        Example::
+
+            s = Search.from_dict({
+                "query": {
+                    "bool": {
+                        "must": [...]
+                    }
+                },
+                "aggs": {...}
+            })
+            s = s.filter('term', published=True)
+        """
+        s = cls()
+        s.update_from_dict(d)
+        return s
+
+    def _clone(self) -> Self:
+        """
+        Return a clone of the current search request. Performs a shallow copy
+        of all the underlying objects. Used internally by most state modifying
+        APIs.
+        """
+        s = super()._clone()
+
+        s._response_class = self._response_class
+        s._knn = [knn.copy() for knn in self._knn]
+        s._rank = self._rank.copy()
+        s._collapse = self._collapse.copy()
+        s._sort = self._sort[:]
+        s._source = copy.copy(self._source) if self._source is not None else None
+        s._highlight = self._highlight.copy()
+        s._highlight_opts = self._highlight_opts.copy()
+        s._suggest = self._suggest.copy()
+        s._script_fields = self._script_fields.copy()
+        for x in ("query", "post_filter"):
+            getattr(s, x)._proxied = getattr(self, x)._proxied
+
+        # copy top-level bucket definitions
+        if self.aggs._params.get("aggs"):
+            s.aggs._params = {"aggs": self.aggs._params["aggs"].copy()}
+        return s
+
+    def response_class(self, cls: Type[Response[_R]]) -> Self:
+        """
+        Override the default wrapper used for the response.
+        """
+        s = self._clone()
+        s._response_class = cls
+        return s
+
+    def update_from_dict(self, d: Dict[str, Any]) -> Self:
+        """
+        Apply options from a serialized body to the current instance. Modifies
+        the object in-place. Used mostly by ``from_dict``.
+        """
+        d = d.copy()
+        if "query" in d:
+            self.query._proxied = Q(d.pop("query"))
+        if "post_filter" in d:
+            self.post_filter._proxied = Q(d.pop("post_filter"))
+
+        aggs = d.pop("aggs", d.pop("aggregations", {}))
+        if aggs:
+            self.aggs._params = {
+                "aggs": {name: A(value) for (name, value) in aggs.items()}
+            }
+        if "knn" in d:
+            self._knn = d.pop("knn")
+            if isinstance(self._knn, dict):
+                self._knn = [self._knn]
+        if "rank" in d:
+            self._rank = d.pop("rank")
+        if "collapse" in d:
+            self._collapse = d.pop("collapse")
+        if "sort" in d:
+            self._sort = d.pop("sort")
+        if "_source" in d:
+            self._source = d.pop("_source")
+        if "highlight" in d:
+            high = d.pop("highlight").copy()
+            self._highlight = high.pop("fields")
+            self._highlight_opts = high
+        if "suggest" in d:
+            self._suggest = d.pop("suggest")
+            if "text" in self._suggest:
+                text = self._suggest.pop("text")
+                for s in self._suggest.values():
+                    s.setdefault("text", text)
+        if "script_fields" in d:
+            self._script_fields = d.pop("script_fields")
+        self._extra.update(d)
+        return self
+
+    def script_fields(self, **kwargs: Any) -> Self:
+        """
+        Define script fields to be calculated on hits. See
+        https://www.elastic.co/guide/en/elasticsearch/reference/current/search-request-script-fields.html
+        for more details.
+
+        Example::
+
+            s = Search()
+            s = s.script_fields(times_two="doc['field'].value * 2")
+            s = s.script_fields(
+                times_three={
+                    'script': {
+                        'lang': 'painless',
+                        'source': "doc['field'].value * params.n",
+                        'params': {'n': 3}
+                    }
+                }
+            )
+
+        """
+        s = self._clone()
+        for name in kwargs:
+            if isinstance(kwargs[name], str):
+                kwargs[name] = {"script": kwargs[name]}
+        s._script_fields.update(kwargs)
+        return s
+
+    def knn(
+        self,
+        field: Union[str, "InstrumentedField"],
+        k: int,
+        num_candidates: int,
+        query_vector: Optional[List[float]] = None,
+        query_vector_builder: Optional[Dict[str, Any]] = None,
+        boost: Optional[float] = None,
+        filter: Optional[Query] = None,
+        similarity: Optional[float] = None,
+        inner_hits: Optional[Dict[str, Any]] = None,
+    ) -> Self:
+        """
+        Add a k-nearest neighbor (kNN) search.
+
+        :arg field: the vector field to search against as a string or document class attribute
+        :arg k: number of nearest neighbors to return as top hits
+        :arg num_candidates: number of nearest neighbor candidates to consider per shard
+        :arg query_vector: the vector to search for
+        :arg query_vector_builder: A dictionary indicating how to build a query vector
+        :arg boost: A floating-point boost factor for kNN scores
+        :arg filter: query to filter the documents that can match
+        :arg similarity: the minimum similarity required for a document to be considered a match, as a float value
+        :arg inner_hits: retrieve hits from nested field
+
+        Example::
+
+            s = Search()
+            s = s.knn(field='embedding', k=5, num_candidates=10, query_vector=vector,
+                      filter=Q('term', category='blog')))
+        """
+        s = self._clone()
+        s._knn.append(
+            {
+                "field": str(field),  # str() is for InstrumentedField instances
+                "k": k,
+                "num_candidates": num_candidates,
+            }
+        )
+        if query_vector is None and query_vector_builder is None:
+            raise ValueError("one of query_vector and query_vector_builder is required")
+        if query_vector is not None and query_vector_builder is not None:
+            raise ValueError(
+                "only one of query_vector and query_vector_builder must be given"
+            )
+        if query_vector is not None:
+            s._knn[-1]["query_vector"] = cast(Any, query_vector)
+        if query_vector_builder is not None:
+            s._knn[-1]["query_vector_builder"] = query_vector_builder
+        if boost is not None:
+            s._knn[-1]["boost"] = boost
+        if filter is not None:
+            if isinstance(filter, Query):
+                s._knn[-1]["filter"] = filter.to_dict()
+            else:
+                s._knn[-1]["filter"] = filter
+        if similarity is not None:
+            s._knn[-1]["similarity"] = similarity
+        if inner_hits is not None:
+            s._knn[-1]["inner_hits"] = inner_hits
+        return s
+
+    def rank(self, rrf: Optional[Union[bool, Dict[str, Any]]] = None) -> Self:
+        """
+        Defines a method for combining and ranking results sets from a combination
+        of searches. Requires a minimum of 2 results sets.
+
+        :arg rrf: Set to ``True`` or an options dictionary to set the rank method to reciprocal rank fusion (RRF).
+
+        Example::
+
+            s = Search()
+            s = s.query('match', content='search text')
+            s = s.knn(field='embedding', k=5, num_candidates=10, query_vector=vector)
+            s = s.rank(rrf=True)
+
+        Note: This option is in technical preview and may change in the future. The syntax will likely change before GA.
+        """
+        s = self._clone()
+        s._rank = {}
+        if rrf is not None and rrf is not False:
+            s._rank["rrf"] = {} if rrf is True else rrf
+        return s
+
+    def source(
+        self,
+        fields: Optional[
+            Union[
+                bool,
+                str,
+                "InstrumentedField",
+                List[Union[str, "InstrumentedField"]],
+                Dict[str, List[Union[str, "InstrumentedField"]]],
+            ]
+        ] = None,
+        **kwargs: Any,
+    ) -> Self:
+        """
+        Selectively control how the _source field is returned.
+
+        :arg fields: field name, wildcard string, list of field names or wildcards,
+                     or dictionary of includes and excludes
+        :arg kwargs: ``includes`` or ``excludes`` arguments, when ``fields`` is ``None``.
+
+        When no arguments are given, the entire document will be returned for
+        each hit.  If ``fields`` is a string or list of strings, the field names or field
+        wildcards given will be included. If ``fields`` is a dictionary with keys of
+        'includes' and/or 'excludes' the fields will be either included or excluded
+        appropriately.
+
+        Calling this multiple times with the same named parameter will override the
+        previous values with the new ones.
+
+        Example::
+
+            s = Search()
+            s = s.source(includes=['obj1.*'], excludes=["*.description"])
+
+            s = Search()
+            s = s.source(includes=['obj1.*']).source(excludes=["*.description"])
+
+        """
+        s = self._clone()
+
+        if fields and kwargs:
+            raise ValueError("You cannot specify fields and kwargs at the same time.")
+
+        @overload
+        def ensure_strings(fields: str) -> str: ...
+
+        @overload
+        def ensure_strings(fields: "InstrumentedField") -> str: ...
+
+        @overload
+        def ensure_strings(
+            fields: List[Union[str, "InstrumentedField"]],
+        ) -> List[str]: ...
+
+        @overload
+        def ensure_strings(
+            fields: Dict[str, List[Union[str, "InstrumentedField"]]],
+        ) -> Dict[str, List[str]]: ...
+
+        def ensure_strings(
+            fields: Union[
+                str,
+                "InstrumentedField",
+                List[Union[str, "InstrumentedField"]],
+                Dict[str, List[Union[str, "InstrumentedField"]]],
+            ],
+        ) -> Union[str, List[str], Dict[str, List[str]]]:
+            if isinstance(fields, dict):
+                return {k: ensure_strings(v) for k, v in fields.items()}
+            elif not isinstance(fields, (str, InstrumentedField)):
+                # we assume that if `fields` is not a any of [dict, str,
+                # InstrumentedField] then it is an iterable of strings or
+                # InstrumentedFields, so we convert them to a plain list of
+                # strings
+                return [str(f) for f in fields]
+            else:
+                return str(fields)
+
+        if fields is not None:
+            s._source = fields if isinstance(fields, bool) else ensure_strings(fields)  # type: ignore[assignment]
+            return s
+
+        if kwargs and not isinstance(s._source, dict):
+            s._source = {}
+
+        if isinstance(s._source, dict):
+            for key, value in kwargs.items():
+                if value is None:
+                    try:
+                        del s._source[key]
+                    except KeyError:
+                        pass
+                else:
+                    s._source[key] = ensure_strings(value)
+
+        return s
+
+    def sort(
+        self, *keys: Union[str, "InstrumentedField", Dict[str, Dict[str, str]]]
+    ) -> Self:
+        """
+        Add sorting information to the search request. If called without
+        arguments it will remove all sort requirements. Otherwise it will
+        replace them. Acceptable arguments are::
+
+            'some.field'
+            '-some.other.field'
+            {'different.field': {'any': 'dict'}}
+
+        so for example::
+
+            s = Search().sort(
+                'category',
+                '-title',
+                {"price" : {"order" : "asc", "mode" : "avg"}}
+            )
+
+        will sort by ``category``, ``title`` (in descending order) and
+        ``price`` in ascending order using the ``avg`` mode.
+
+        The API returns a copy of the Search object and can thus be chained.
+        """
+        s = self._clone()
+        s._sort = []
+        for k in keys:
+            if not isinstance(k, dict):
+                sort_field = str(k)
+                if sort_field.startswith("-"):
+                    if sort_field[1:] == "_score":
+                        raise IllegalOperation("Sorting by `-_score` is not allowed.")
+                    s._sort.append({sort_field[1:]: {"order": "desc"}})
+                else:
+                    s._sort.append(sort_field)
+            else:
+                s._sort.append(k)
+        return s
+
+    def collapse(
+        self,
+        field: Optional[Union[str, "InstrumentedField"]] = None,
+        inner_hits: Optional[Dict[str, Any]] = None,
+        max_concurrent_group_searches: Optional[int] = None,
+    ) -> Self:
+        """
+        Add collapsing information to the search request.
+        If called without providing ``field``, it will remove all collapse
+        requirements, otherwise it will replace them with the provided
+        arguments.
+        The API returns a copy of the Search object and can thus be chained.
+        """
+        s = self._clone()
+        s._collapse = {}
+
+        if field is None:
+            return s
+
+        s._collapse["field"] = str(field)
+        if inner_hits:
+            s._collapse["inner_hits"] = inner_hits
+        if max_concurrent_group_searches:
+            s._collapse["max_concurrent_group_searches"] = max_concurrent_group_searches
+        return s
+
+    def highlight_options(self, **kwargs: Any) -> Self:
+        """
+        Update the global highlighting options used for this request. For
+        example::
+
+            s = Search()
+            s = s.highlight_options(order='score')
+        """
+        s = self._clone()
+        s._highlight_opts.update(kwargs)
+        return s
+
+    def highlight(
+        self, *fields: Union[str, "InstrumentedField"], **kwargs: Any
+    ) -> Self:
+        """
+        Request highlighting of some fields. All keyword arguments passed in will be
+        used as parameters for all the fields in the ``fields`` parameter. Example::
+
+            Search().highlight('title', 'body', fragment_size=50)
+
+        will produce the equivalent of::
+
+            {
+                "highlight": {
+                    "fields": {
+                        "body": {"fragment_size": 50},
+                        "title": {"fragment_size": 50}
+                    }
+                }
+            }
+
+        If you want to have different options for different fields
+        you can call ``highlight`` twice::
+
+            Search().highlight('title', fragment_size=50).highlight('body', fragment_size=100)
+
+        which will produce::
+
+            {
+                "highlight": {
+                    "fields": {
+                        "body": {"fragment_size": 100},
+                        "title": {"fragment_size": 50}
+                    }
+                }
+            }
+
+        """
+        s = self._clone()
+        for f in fields:
+            s._highlight[str(f)] = kwargs
+        return s
+
+    def suggest(
+        self,
+        name: str,
+        text: Optional[str] = None,
+        regex: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Self:
+        """
+        Add a suggestions request to the search.
+
+        :arg name: name of the suggestion
+        :arg text: text to suggest on
+
+        All keyword arguments will be added to the suggestions body. For example::
+
+            s = Search()
+            s = s.suggest('suggestion-1', 'Elasticsearch', term={'field': 'body'})
+
+        # regex query for Completion Suggester
+            s = Search()
+            s = s.suggest('suggestion-1', regex='py[thon|py]', completion={'field': 'body'})
+        """
+        if text is None and regex is None:
+            raise ValueError('You have to pass "text" or "regex" argument.')
+        if text and regex:
+            raise ValueError('You can only pass either "text" or "regex" argument.')
+        if regex and "completion" not in kwargs:
+            raise ValueError(
+                '"regex" argument must be passed with "completion" keyword argument.'
+            )
+
+        s = self._clone()
+        if regex:
+            s._suggest[name] = {"regex": regex}
+        elif text:
+            if "completion" in kwargs:
+                s._suggest[name] = {"prefix": text}
+            else:
+                s._suggest[name] = {"text": text}
+        s._suggest[name].update(kwargs)
+        return s
+
+    def search_after(self) -> Self:
+        """
+        Return a ``Search`` instance that retrieves the next page of results.
+
+        This method provides an easy way to paginate a long list of results using
+        the ``search_after`` option. For example::
+
+            page_size = 20
+            s = Search()[:page_size].sort("date")
+
+            while True:
+                # get a page of results
+                r = await s.execute()
+
+                # do something with this page of results
+
+                # exit the loop if we reached the end
+                if len(r.hits) < page_size:
+                    break
+
+                # get a search object with the next page of results
+                s = s.search_after()
+
+        Note that the ``search_after`` option requires the search to have an
+        explicit ``sort`` order.
+        """
+        if not hasattr(self, "_response"):
+            raise ValueError("A search must be executed before using search_after")
+        return cast(Self, self._response.search_after())
+
+    def to_dict(self, count: bool = False, **kwargs: Any) -> Dict[str, Any]:
+        """
+        Serialize the search into the dictionary that will be sent over as the
+        request's body.
+
+        :arg count: a flag to specify if we are interested in a body for count -
+            no aggregations, no pagination bounds etc.
+
+        All additional keyword arguments will be included into the dictionary.
+        """
+        d = {}
+
+        if self.query:
+            d["query"] = recursive_to_dict(self.query)
+
+        if self._knn:
+            if len(self._knn) == 1:
+                d["knn"] = self._knn[0]
+            else:
+                d["knn"] = self._knn
+
+        if self._rank:
+            d["rank"] = self._rank
+
+        # count request doesn't care for sorting and other things
+        if not count:
+            if self.post_filter:
+                d["post_filter"] = recursive_to_dict(self.post_filter.to_dict())
+
+            if self.aggs.aggs:
+                d.update(recursive_to_dict(self.aggs.to_dict()))
+
+            if self._sort:
+                d["sort"] = self._sort
+
+            if self._collapse:
+                d["collapse"] = self._collapse
+
+            d.update(recursive_to_dict(self._extra))
+
+            if self._source not in (None, {}):
+                d["_source"] = self._source
+
+            if self._highlight:
+                d["highlight"] = {"fields": self._highlight}
+                d["highlight"].update(self._highlight_opts)
+
+            if self._suggest:
+                d["suggest"] = self._suggest
+
+            if self._script_fields:
+                d["script_fields"] = self._script_fields
+
+        d.update(recursive_to_dict(kwargs))
+        return d
+
+
+class MultiSearchBase(Request[_R]):
+    """
+    Combine multiple :class:`~elasticsearch.dsl.Search` objects into a single
+    request.
+    """
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+        self._searches: List[SearchBase[_R]] = []
+
+    def __getitem__(self, key: Union[int, slice]) -> Any:
+        return self._searches[key]
+
+    def __iter__(self) -> Iterator[SearchBase[_R]]:
+        return iter(self._searches)
+
+    def _clone(self) -> Self:
+        ms = super()._clone()
+        ms._searches = self._searches[:]
+        return ms
+
+    def add(self, search: SearchBase[_R]) -> Self:
+        """
+        Adds a new :class:`~elasticsearch.dsl.Search` object to the request::
+
+            ms = MultiSearch(index='my-index')
+            ms = ms.add(Search(doc_type=Category).filter('term', category='python'))
+            ms = ms.add(Search(doc_type=Blog))
+        """
+        ms = self._clone()
+        ms._searches.append(search)
+        return ms
+
+    def to_dict(self) -> List[Dict[str, Any]]:  # type: ignore[override]
+        out: List[Dict[str, Any]] = []
+        for s in self._searches:
+            meta: Dict[str, Any] = {}
+            if s._index:
+                meta["index"] = cast(Any, s._index)
+            meta.update(s._params)
+
+            out.append(meta)
+            out.append(s.to_dict())
+
+        return out