elasticsearch9 9.1.3__py3-none-any.whl → 9.2.1__py3-none-any.whl

elasticsearch9/dsl/pydantic.py

@@ -0,0 +1,152 @@
+#  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.
+
+from typing import Any, ClassVar, Dict, List, Optional, Tuple, Type
+
+from pydantic import BaseModel, Field, PrivateAttr
+from typing_extensions import Annotated, Self, dataclass_transform
+
+from .. import dsl
+
+
+class ESMeta(BaseModel):
+    """Metadata items associated with Elasticsearch documents."""
+
+    id: str = ""
+    index: str = ""
+    primary_term: int = 0
+    seq_no: int = 0
+    version: int = 0
+    score: float = 0
+
+
+class _BaseModel(BaseModel):
+    meta: Annotated[ESMeta, dsl.mapped_field(exclude=True)] = Field(
+        default=ESMeta(),
+        init=False,
+    )
+
+
+class _BaseESModelMetaclass(type(BaseModel)):  # type: ignore[misc]
+    """Generic metaclass methods for BaseEsModel and AsyncBaseESModel."""
+
+    @staticmethod
+    def process_annotations(
+        metacls: Type["_BaseESModelMetaclass"], annotations: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Process Pydantic typing annotations and adapt them so that they can
+        be used to create the Elasticsearch document.
+        """
+        updated_annotations = {}
+        for var, ann in annotations.items():
+            if isinstance(ann, type(BaseModel)):
+                # an inner Pydantic model is transformed into an Object field
+                updated_annotations[var] = metacls.make_dsl_class(
+                    metacls, dsl.InnerDoc, ann
+                )
+            elif (
+                hasattr(ann, "__origin__")
+                and ann.__origin__ in [list, List]
+                and isinstance(ann.__args__[0], type(BaseModel))
+            ):
+                # an inner list of Pydantic models is transformed into a Nested field
+                updated_annotations[var] = List[  # type: ignore[assignment,misc]
+                    metacls.make_dsl_class(metacls, dsl.InnerDoc, ann.__args__[0])
+                ]
+            else:
+                updated_annotations[var] = ann
+        return updated_annotations
+
+    @staticmethod
+    def make_dsl_class(
+        metacls: Type["_BaseESModelMetaclass"],
+        dsl_class: type,
+        pydantic_model: type,
+        pydantic_attrs: Optional[Dict[str, Any]] = None,
+    ) -> type:
+        """Create a DSL document class dynamically, using the structure of a
+        Pydantic model."""
+        dsl_attrs = {
+            attr: value
+            for attr, value in dsl_class.__dict__.items()
+            if not attr.startswith("__")
+        }
+        pydantic_attrs = {
+            **(pydantic_attrs or {}),
+            "__annotations__": metacls.process_annotations(
+                metacls, pydantic_model.__annotations__
+            ),
+        }
+        return type(dsl_class)(
+            f"_ES{pydantic_model.__name__}",
+            (dsl_class,),
+            {
+                **pydantic_attrs,
+                **dsl_attrs,
+                "__qualname__": f"_ES{pydantic_model.__name__}",
+            },
+        )
+
+
+class BaseESModelMetaclass(_BaseESModelMetaclass):
+    """Metaclass for the BaseESModel class."""
+
+    def __new__(cls, name: str, bases: Tuple[type, ...], attrs: Dict[str, Any]) -> Any:
+        model = super().__new__(cls, name, bases, attrs)
+        model._doc = cls.make_dsl_class(cls, dsl.Document, model, attrs)
+        return model
+
+
+class AsyncBaseESModelMetaclass(_BaseESModelMetaclass):
+    """Metaclass for the AsyncBaseESModel class."""
+
+    def __new__(cls, name: str, bases: Tuple[type, ...], attrs: Dict[str, Any]) -> Any:
+        model = super().__new__(cls, name, bases, attrs)
+        model._doc = cls.make_dsl_class(cls, dsl.AsyncDocument, model, attrs)
+        return model
+
+
+@dataclass_transform(kw_only_default=True, field_specifiers=(Field, PrivateAttr))
+class BaseESModel(_BaseModel, metaclass=BaseESModelMetaclass):
+    _doc: ClassVar[Type[dsl.Document]]
+
+    def to_doc(self) -> dsl.Document:
+        """Convert this model to an Elasticsearch document."""
+        data = self.model_dump()
+        meta = {f"_{k}": v for k, v in data.pop("meta", {}).items() if v}
+        return self._doc(**meta, **data)
+
+    @classmethod
+    def from_doc(cls, dsl_obj: dsl.Document) -> Self:
+        """Create a model from the given Elasticsearch document."""
+        return cls(meta=ESMeta(**dsl_obj.meta.to_dict()), **dsl_obj.to_dict())
+
+
+@dataclass_transform(kw_only_default=True, field_specifiers=(Field, PrivateAttr))
+class AsyncBaseESModel(_BaseModel, metaclass=AsyncBaseESModelMetaclass):
+    _doc: ClassVar[Type[dsl.AsyncDocument]]
+
+    def to_doc(self) -> dsl.AsyncDocument:
+        """Convert this model to an Elasticsearch document."""
+        data = self.model_dump()
+        meta = {f"_{k}": v for k, v in data.pop("meta", {}).items() if v}
+        return self._doc(**meta, **data)
+
+    @classmethod
+    def from_doc(cls, dsl_obj: dsl.AsyncDocument) -> Self:
+        """Create a model from the given Elasticsearch document."""
+        return cls(meta=ESMeta(**dsl_obj.meta.to_dict()), **dsl_obj.to_dict())

elasticsearch9/dsl/query.py

@@ -1087,6 +1087,8 @@ class Knn(Query):
         a query_vector_builder or query_vector, but not both.
     :arg num_candidates: The number of nearest neighbor candidates to
         consider per shard
+    :arg visit_percentage: The percentage of vectors to explore per shard
+        while doing knn search with bbq_disk
     :arg k: The final number of nearest neighbors to return as top hits
     :arg filter: Filters for the kNN search query
     :arg similarity: The minimum similarity for a vector to be considered
@@ -1115,6 +1117,7 @@ class Knn(Query):
             "types.QueryVectorBuilder", Dict[str, Any], "DefaultType"
         ] = DEFAULT,
         num_candidates: Union[int, "DefaultType"] = DEFAULT,
+        visit_percentage: Union[float, "DefaultType"] = DEFAULT,
         k: Union[int, "DefaultType"] = DEFAULT,
         filter: Union[Query, Sequence[Query], "DefaultType"] = DEFAULT,
         similarity: Union[float, "DefaultType"] = DEFAULT,
@@ -1130,6 +1133,7 @@ class Knn(Query):
             query_vector=query_vector,
             query_vector_builder=query_vector_builder,
             num_candidates=num_candidates,
+            visit_percentage=visit_percentage,
             k=k,
             filter=filter,
             similarity=similarity,
@@ -1445,7 +1449,7 @@ class MoreLikeThis(Query):
         ] = DEFAULT,
         version: Union[int, "DefaultType"] = DEFAULT,
         version_type: Union[
-            Literal["internal", "external", "external_gte", "force"], "DefaultType"
+            Literal["internal", "external", "external_gte"], "DefaultType"
         ] = DEFAULT,
         boost: Union[float, "DefaultType"] = DEFAULT,
         _name: Union[str, "DefaultType"] = DEFAULT,

elasticsearch9/dsl/search_base.py

@@ -721,14 +721,18 @@ class SearchBase(Request[_R]):
 
         def ensure_strings(
             fields: Union[
+                bool,
                 str,
                 "InstrumentedField",
                 List[Union[str, "InstrumentedField"]],
                 Dict[str, List[Union[str, "InstrumentedField"]]],
             ],
-        ) -> Union[str, List[str], Dict[str, List[str]]]:
+        ) -> Union[bool, str, List[str], Dict[str, List[str]]]:
             if isinstance(fields, dict):
                 return {k: ensure_strings(v) for k, v in fields.items()}
+            elif isinstance(fields, bool):
+                # boolean settings should stay the way they are
+                return fields
             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

elasticsearch9/dsl/types.py

@@ -151,9 +151,10 @@ class ChunkingSettings(AttrDict[Any]):
         strategies in the linked documentation. Defaults to `sentence` if
         omitted.
     :arg max_chunk_size: (required) The maximum size of a chunk in words.
-        This value cannot be higher than `300` or lower than `20` (for
-        `sentence` strategy) or `10` (for `word` strategy). Defaults to
-        `250` if omitted.
+        This value cannot be lower than `20` (for `sentence` strategy) or
+        `10` (for `word` strategy). This value should not exceed the
+        window size for the associated model. Defaults to `250` if
+        omitted.
     :arg separator_group: Only applicable to the `recursive` strategy and
         required when using it.  Sets a predefined list of separators in
         the saved chunking settings based on the selected text type.
@@ -397,14 +398,15 @@ class DenseVectorIndexOptions(AttrDict[Any]):
         HNSW graph.  Only applicable to `hnsw`, `int8_hnsw`, `bbq_hnsw`,
         and `int4_hnsw` index types. Defaults to `16` if omitted.
     :arg rescore_vector: The rescore vector options. This is only
-        applicable to `bbq_hnsw`, `int4_hnsw`, `int8_hnsw`, `bbq_flat`,
-        `int4_flat`, and `int8_flat` index types.
+        applicable to `bbq_disk`, `bbq_hnsw`, `int4_hnsw`, `int8_hnsw`,
+        `bbq_flat`, `int4_flat`, and `int8_flat` index types.
     """
 
     type: Union[
         Literal[
             "bbq_flat",
             "bbq_hnsw",
+            "bbq_disk",
             "flat",
             "hnsw",
             "int4_flat",
@@ -428,6 +430,7 @@ class DenseVectorIndexOptions(AttrDict[Any]):
             Literal[
                 "bbq_flat",
                 "bbq_hnsw",
+                "bbq_disk",
                 "flat",
                 "hnsw",
                 "int4_flat",
@@ -2327,9 +2330,7 @@ class LikeDocument(AttrDict[Any]):
     per_field_analyzer: Union[Mapping[Union[str, InstrumentedField], str], DefaultType]
     routing: Union[str, DefaultType]
     version: Union[int, DefaultType]
-    version_type: Union[
-        Literal["internal", "external", "external_gte", "force"], DefaultType
-    ]
+    version_type: Union[Literal["internal", "external", "external_gte"], DefaultType]
 
     def __init__(
         self,
@@ -2344,7 +2345,7 @@ class LikeDocument(AttrDict[Any]):
         routing: Union[str, DefaultType] = DEFAULT,
         version: Union[int, DefaultType] = DEFAULT,
         version_type: Union[
-            Literal["internal", "external", "external_gte", "force"], DefaultType
+            Literal["internal", "external", "external_gte"], DefaultType
         ] = DEFAULT,
         **kwargs: Any,
     ):
@@ -3190,6 +3191,33 @@ class ScriptedHeuristic(AttrDict[Any]):
         super().__init__(kwargs)
 
 
+class SemanticTextIndexOptions(AttrDict[Any]):
+    """
+    :arg dense_vector:
+    :arg sparse_vector:
+    """
+
+    dense_vector: Union["DenseVectorIndexOptions", Dict[str, Any], DefaultType]
+    sparse_vector: Union["SparseVectorIndexOptions", Dict[str, Any], DefaultType]
+
+    def __init__(
+        self,
+        *,
+        dense_vector: Union[
+            "DenseVectorIndexOptions", Dict[str, Any], DefaultType
+        ] = DEFAULT,
+        sparse_vector: Union[
+            "SparseVectorIndexOptions", Dict[str, Any], DefaultType
+        ] = DEFAULT,
+        **kwargs: Any,
+    ):
+        if dense_vector is not DEFAULT:
+            kwargs["dense_vector"] = dense_vector
+        if sparse_vector is not DEFAULT:
+            kwargs["sparse_vector"] = sparse_vector
+        super().__init__(kwargs)
+
+
 class ShapeFieldQuery(AttrDict[Any]):
     """
     :arg indexed_shape: Queries using a pre-indexed shape.