elasticsearch 8.14.0__py3-none-any.whl → 8.15.1__py3-none-any.whl

elasticsearch/helpers/actions.py

@@ -34,6 +34,8 @@ from typing import (
     Union,
 )
 
+from elastic_transport import OpenTelemetrySpan
+
 from .. import Elasticsearch
 from ..compat import to_bytes
 from ..exceptions import ApiError, NotFoundError, TransportError
@@ -322,6 +324,7 @@ def _process_bulk_chunk(
             Tuple[_TYPE_BULK_ACTION_HEADER, _TYPE_BULK_ACTION_BODY],
         ]
     ],
+    otel_span: OpenTelemetrySpan,
     raise_on_exception: bool = True,
     raise_on_error: bool = True,
     ignore_status: Union[int, Collection[int]] = (),
@@ -331,28 +334,29 @@ def _process_bulk_chunk(
     """
     Send a bulk request to elasticsearch and process the output.
     """
-    if isinstance(ignore_status, int):
-        ignore_status = (ignore_status,)
-
-    try:
-        # send the actual request
-        resp = client.bulk(*args, operations=bulk_actions, **kwargs)  # type: ignore[arg-type]
-    except ApiError as e:
-        gen = _process_bulk_chunk_error(
-            error=e,
-            bulk_data=bulk_data,
-            ignore_status=ignore_status,
-            raise_on_exception=raise_on_exception,
-            raise_on_error=raise_on_error,
-        )
-    else:
-        gen = _process_bulk_chunk_success(
-            resp=resp.body,
-            bulk_data=bulk_data,
-            ignore_status=ignore_status,
-            raise_on_error=raise_on_error,
-        )
-    yield from gen
+    with client._otel.use_span(otel_span):
+        if isinstance(ignore_status, int):
+            ignore_status = (ignore_status,)
+
+        try:
+            # send the actual request
+            resp = client.bulk(*args, operations=bulk_actions, **kwargs)  # type: ignore[arg-type]
+        except ApiError as e:
+            gen = _process_bulk_chunk_error(
+                error=e,
+                bulk_data=bulk_data,
+                ignore_status=ignore_status,
+                raise_on_exception=raise_on_exception,
+                raise_on_error=raise_on_error,
+            )
+        else:
+            gen = _process_bulk_chunk_success(
+                resp=resp.body,
+                bulk_data=bulk_data,
+                ignore_status=ignore_status,
+                raise_on_error=raise_on_error,
+            )
+        yield from gen
 
 
 def streaming_bulk(
@@ -370,6 +374,7 @@ def streaming_bulk(
     max_backoff: float = 600,
     yield_ok: bool = True,
     ignore_status: Union[int, Collection[int]] = (),
+    span_name: str = "helpers.streaming_bulk",
     *args: Any,
     **kwargs: Any,
 ) -> Iterable[Tuple[bool, Dict[str, Any]]]:
@@ -406,73 +411,78 @@ def streaming_bulk(
     :arg yield_ok: if set to False will skip successful documents in the output
     :arg ignore_status: list of HTTP status code that you want to ignore
     """
-    client = client.options()
-    client._client_meta = (("h", "bp"),)
+    with client._otel.helpers_span(span_name) as otel_span:
+        client = client.options()
+        client._client_meta = (("h", "bp"),)
 
-    serializer = client.transport.serializers.get_serializer("application/json")
+        serializer = client.transport.serializers.get_serializer("application/json")
 
-    bulk_data: List[
-        Union[
-            Tuple[_TYPE_BULK_ACTION_HEADER],
-            Tuple[_TYPE_BULK_ACTION_HEADER, _TYPE_BULK_ACTION_BODY],
+        bulk_data: List[
+            Union[
+                Tuple[_TYPE_BULK_ACTION_HEADER],
+                Tuple[_TYPE_BULK_ACTION_HEADER, _TYPE_BULK_ACTION_BODY],
+            ]
         ]
-    ]
-    bulk_actions: List[bytes]
-    for bulk_data, bulk_actions in _chunk_actions(
-        map(expand_action_callback, actions), chunk_size, max_chunk_bytes, serializer
-    ):
-        for attempt in range(max_retries + 1):
-            to_retry: List[bytes] = []
-            to_retry_data: List[
-                Union[
-                    Tuple[_TYPE_BULK_ACTION_HEADER],
-                    Tuple[_TYPE_BULK_ACTION_HEADER, _TYPE_BULK_ACTION_BODY],
-                ]
-            ] = []
-            if attempt:
-                time.sleep(min(max_backoff, initial_backoff * 2 ** (attempt - 1)))
-
-            try:
-                for data, (ok, info) in zip(
-                    bulk_data,
-                    _process_bulk_chunk(
-                        client,
-                        bulk_actions,
+        bulk_actions: List[bytes]
+        for bulk_data, bulk_actions in _chunk_actions(
+            map(expand_action_callback, actions),
+            chunk_size,
+            max_chunk_bytes,
+            serializer,
+        ):
+            for attempt in range(max_retries + 1):
+                to_retry: List[bytes] = []
+                to_retry_data: List[
+                    Union[
+                        Tuple[_TYPE_BULK_ACTION_HEADER],
+                        Tuple[_TYPE_BULK_ACTION_HEADER, _TYPE_BULK_ACTION_BODY],
+                    ]
+                ] = []
+                if attempt:
+                    time.sleep(min(max_backoff, initial_backoff * 2 ** (attempt - 1)))
+
+                try:
+                    for data, (ok, info) in zip(
                         bulk_data,
-                        raise_on_exception,
-                        raise_on_error,
-                        ignore_status,
-                        *args,
-                        **kwargs,
-                    ),
-                ):
-                    if not ok:
-                        action, info = info.popitem()
-                        # retry if retries enabled, we get 429, and we are not
-                        # in the last attempt
-                        if (
-                            max_retries
-                            and info["status"] == 429
-                            and (attempt + 1) <= max_retries
-                        ):
-                            # _process_bulk_chunk expects bytes so we need to
-                            # re-serialize the data
-                            to_retry.extend(map(serializer.dumps, data))
-                            to_retry_data.append(data)
-                        else:
-                            yield ok, {action: info}
-                    elif yield_ok:
-                        yield ok, info
-
-            except ApiError as e:
-                # suppress 429 errors since we will retry them
-                if attempt == max_retries or e.status_code != 429:
-                    raise
-            else:
-                if not to_retry:
-                    break
-                # retry only subset of documents that didn't succeed
-                bulk_actions, bulk_data = to_retry, to_retry_data
+                        _process_bulk_chunk(
+                            client,
+                            bulk_actions,
+                            bulk_data,
+                            otel_span,
+                            raise_on_exception,
+                            raise_on_error,
+                            ignore_status,
+                            *args,
+                            **kwargs,
+                        ),
+                    ):
+                        if not ok:
+                            action, info = info.popitem()
+                            # retry if retries enabled, we get 429, and we are not
+                            # in the last attempt
+                            if (
+                                max_retries
+                                and info["status"] == 429
+                                and (attempt + 1) <= max_retries
+                            ):
+                                # _process_bulk_chunk expects bytes so we need to
+                                # re-serialize the data
+                                to_retry.extend(map(serializer.dumps, data))
+                                to_retry_data.append(data)
+                            else:
+                                yield ok, {action: info}
+                        elif yield_ok:
+                            yield ok, info
+
+                except ApiError as e:
+                    # suppress 429 errors since we will retry them
+                    if attempt == max_retries or e.status_code != 429:
+                        raise
+                else:
+                    if not to_retry:
+                        break
+                    # retry only subset of documents that didn't succeed
+                    bulk_actions, bulk_data = to_retry, to_retry_data
 
 
 def bulk(
@@ -519,7 +529,7 @@ def bulk(
     # make streaming_bulk yield successful results so we can count them
     kwargs["yield_ok"] = True
     for ok, item in streaming_bulk(
-        client, actions, ignore_status=ignore_status, *args, **kwargs  # type: ignore[misc]
+        client, actions, ignore_status=ignore_status, span_name="helpers.bulk", *args, **kwargs  # type: ignore[misc]
     ):
         # go through request-response pairs and detect failures
         if not ok:
@@ -589,27 +599,31 @@ def parallel_bulk(
             ] = Queue(max(queue_size, thread_count))
             self._quick_put = self._inqueue.put
 
-    pool = BlockingPool(thread_count)
+    with client._otel.helpers_span("helpers.parallel_bulk") as otel_span:
+        pool = BlockingPool(thread_count)
 
-    try:
-        for result in pool.imap(
-            lambda bulk_chunk: list(
-                _process_bulk_chunk(
-                    client,
-                    bulk_chunk[1],
-                    bulk_chunk[0],
-                    ignore_status=ignore_status,  # type: ignore[misc]
-                    *args,
-                    **kwargs,
-                )
-            ),
-            _chunk_actions(expanded_actions, chunk_size, max_chunk_bytes, serializer),
-        ):
-            yield from result
-
-    finally:
-        pool.close()
-        pool.join()
+        try:
+            for result in pool.imap(
+                lambda bulk_chunk: list(
+                    _process_bulk_chunk(
+                        client,
+                        bulk_chunk[1],
+                        bulk_chunk[0],
+                        otel_span=otel_span,
+                        ignore_status=ignore_status,  # type: ignore[misc]
+                        *args,
+                        **kwargs,
+                    )
+                ),
+                _chunk_actions(
+                    expanded_actions, chunk_size, max_chunk_bytes, serializer
+                ),
+            ):
+                yield from result
+
+        finally:
+            pool.close()
+            pool.join()
 
 
 def scan(

elasticsearch/helpers/vectorstore/_async/vectorstore.py

@@ -60,6 +60,7 @@ class AsyncVectorStore:
         vector_field: str = "vector_field",
         metadata_mappings: Optional[Dict[str, Any]] = None,
         user_agent: str = f"elasticsearch-py-vs/{lib_version}",
+        custom_index_settings: Optional[Dict[str, Any]] = None,
     ) -> None:
         """
         :param user_header: user agent header specific to the 3rd party integration.
@@ -72,6 +73,11 @@ class AsyncVectorStore:
             the embedding vector goes in this field.
         :param client: Elasticsearch client connection. Alternatively specify the
             Elasticsearch connection with the other es_* parameters.
+        :param custom_index_settings: A dictionary of custom settings for the index.
+            This can include configurations like the number of shards, number of replicas,
+            analysis settings, and other index-specific settings. If not provided, default
+            settings will be used. Note that if the same setting is provided by both the user
+            and the strategy, will raise an error.
         """
         # Add integration-specific usage header for tracking usage in Elastic Cloud.
         # client.options preserves existing (non-user-agent) headers.
@@ -90,6 +96,7 @@ class AsyncVectorStore:
         self.text_field = text_field
         self.vector_field = vector_field
         self.metadata_mappings = metadata_mappings
+        self.custom_index_settings = custom_index_settings
 
     async def close(self) -> None:
         return await self.client.close()
@@ -225,7 +232,7 @@ class AsyncVectorStore:
     async def search(
         self,
         *,
-        query: Optional[str],
+        query: Optional[str] = None,
         query_vector: Optional[List[float]] = None,
         k: int = 4,
         num_candidates: int = 50,
@@ -306,6 +313,16 @@ class AsyncVectorStore:
             vector_field=self.vector_field,
             num_dimensions=self.num_dimensions,
         )
+
+        if self.custom_index_settings:
+            conflicting_keys = set(self.custom_index_settings.keys()) & set(
+                settings.keys()
+            )
+            if conflicting_keys:
+                raise ValueError(f"Conflicting settings: {conflicting_keys}")
+            else:
+                settings.update(self.custom_index_settings)
+
         if self.metadata_mappings:
             metadata = mappings["properties"].get("metadata", {"properties": {}})
             for key in self.metadata_mappings.keys():
@@ -327,8 +344,9 @@ class AsyncVectorStore:
     async def max_marginal_relevance_search(
         self,
         *,
-        embedding_service: AsyncEmbeddingService,
-        query: str,
+        query: Optional[str] = None,
+        query_embedding: Optional[List[float]] = None,
+        embedding_service: Optional[AsyncEmbeddingService] = None,
         vector_field: str,
         k: int = 4,
         num_candidates: int = 20,
@@ -344,6 +362,8 @@ class AsyncVectorStore:
             among selected documents.
 
         :param query (str): Text to look up documents similar to.
+        :param query_embedding: Input embedding vector. If given, input query string is
+            ignored.
         :param k (int): Number of Documents to return. Defaults to 4.
         :param fetch_k (int): Number of Documents to fetch to pass to MMR algorithm.
         :param lambda_mult (float): Number between 0 and 1 that determines the degree
@@ -364,12 +384,22 @@ class AsyncVectorStore:
             remove_vector_query_field_from_metadata = False
 
         # Embed the query
-        query_embedding = await embedding_service.embed_query(query)
+        if query_embedding:
+            query_vector = query_embedding
+        else:
+            if not query:
+                raise ValueError("specify either query or query_embedding to search")
+            elif embedding_service:
+                query_vector = await embedding_service.embed_query(query)
+            elif self.embedding_service:
+                query_vector = await self.embedding_service.embed_query(query)
+            else:
+                raise ValueError("specify embedding_service to search with query")
 
         # Fetch the initial documents
         got_hits = await self.search(
             query=None,
-            query_vector=query_embedding,
+            query_vector=query_vector,
             k=num_candidates,
             fields=fields,
             custom_query=custom_query,
@@ -380,7 +410,7 @@ class AsyncVectorStore:
 
         # Select documents using maximal marginal relevance
         selected_indices = maximal_marginal_relevance(
-            query_embedding, got_embeddings, lambda_mult=lambda_mult, k=k
+            query_vector, got_embeddings, lambda_mult=lambda_mult, k=k
         )
         selected_hits = [got_hits[i] for i in selected_indices]
 

elasticsearch/helpers/vectorstore/_sync/vectorstore.py

@@ -57,6 +57,7 @@ class VectorStore:
         vector_field: str = "vector_field",
         metadata_mappings: Optional[Dict[str, Any]] = None,
         user_agent: str = f"elasticsearch-py-vs/{lib_version}",
+        custom_index_settings: Optional[Dict[str, Any]] = None,
     ) -> None:
         """
         :param user_header: user agent header specific to the 3rd party integration.
@@ -69,6 +70,11 @@ class VectorStore:
             the embedding vector goes in this field.
         :param client: Elasticsearch client connection. Alternatively specify the
             Elasticsearch connection with the other es_* parameters.
+        :param custom_index_settings: A dictionary of custom settings for the index.
+            This can include configurations like the number of shards, number of replicas,
+            analysis settings, and other index-specific settings. If not provided, default
+            settings will be used. Note that if the same setting is provided by both the user
+            and the strategy, will raise an error.
         """
         # Add integration-specific usage header for tracking usage in Elastic Cloud.
         # client.options preserves existing (non-user-agent) headers.
@@ -87,6 +93,7 @@ class VectorStore:
         self.text_field = text_field
         self.vector_field = vector_field
         self.metadata_mappings = metadata_mappings
+        self.custom_index_settings = custom_index_settings
 
     def close(self) -> None:
         return self.client.close()
@@ -222,7 +229,7 @@ class VectorStore:
     def search(
         self,
         *,
-        query: Optional[str],
+        query: Optional[str] = None,
         query_vector: Optional[List[float]] = None,
         k: int = 4,
         num_candidates: int = 50,
@@ -303,6 +310,16 @@ class VectorStore:
             vector_field=self.vector_field,
             num_dimensions=self.num_dimensions,
         )
+
+        if self.custom_index_settings:
+            conflicting_keys = set(self.custom_index_settings.keys()) & set(
+                settings.keys()
+            )
+            if conflicting_keys:
+                raise ValueError(f"Conflicting settings: {conflicting_keys}")
+            else:
+                settings.update(self.custom_index_settings)
+
         if self.metadata_mappings:
             metadata = mappings["properties"].get("metadata", {"properties": {}})
             for key in self.metadata_mappings.keys():
@@ -324,8 +341,9 @@ class VectorStore:
     def max_marginal_relevance_search(
         self,
         *,
-        embedding_service: EmbeddingService,
-        query: str,
+        query: Optional[str] = None,
+        query_embedding: Optional[List[float]] = None,
+        embedding_service: Optional[EmbeddingService] = None,
         vector_field: str,
         k: int = 4,
         num_candidates: int = 20,
@@ -341,6 +359,8 @@ class VectorStore:
             among selected documents.
 
         :param query (str): Text to look up documents similar to.
+        :param query_embedding: Input embedding vector. If given, input query string is
+            ignored.
         :param k (int): Number of Documents to return. Defaults to 4.
         :param fetch_k (int): Number of Documents to fetch to pass to MMR algorithm.
         :param lambda_mult (float): Number between 0 and 1 that determines the degree
@@ -361,12 +381,22 @@ class VectorStore:
             remove_vector_query_field_from_metadata = False
 
         # Embed the query
-        query_embedding = embedding_service.embed_query(query)
+        if query_embedding:
+            query_vector = query_embedding
+        else:
+            if not query:
+                raise ValueError("specify either query or query_embedding to search")
+            elif embedding_service:
+                query_vector = embedding_service.embed_query(query)
+            elif self.embedding_service:
+                query_vector = self.embedding_service.embed_query(query)
+            else:
+                raise ValueError("specify embedding_service to search with query")
 
         # Fetch the initial documents
         got_hits = self.search(
             query=None,
-            query_vector=query_embedding,
+            query_vector=query_vector,
             k=num_candidates,
             fields=fields,
             custom_query=custom_query,
@@ -377,7 +407,7 @@ class VectorStore:
 
         # Select documents using maximal marginal relevance
         selected_indices = maximal_marginal_relevance(
-            query_embedding, got_embeddings, lambda_mult=lambda_mult, k=k
+            query_vector, got_embeddings, lambda_mult=lambda_mult, k=k
         )
         selected_hits = [got_hits[i] for i in selected_indices]
 

elasticsearch/serializer.py

@@ -49,6 +49,14 @@ except ImportError:
     _OrjsonSerializer = None  # type: ignore[assignment,misc]
 
 
+try:
+    import pyarrow as pa
+
+    __all__.append("PyArrowSerializer")
+except ImportError:
+    pa = None
+
+
 class JsonSerializer(_JsonSerializer):
     mimetype: ClassVar[str] = "application/json"
 
@@ -114,6 +122,29 @@ class MapboxVectorTileSerializer(Serializer):
         raise SerializationError(f"Cannot serialize {data!r} into a MapBox vector tile")
 
 
+if pa is not None:
+
+    class PyArrowSerializer(Serializer):
+        """PyArrow serializer for deserializing Arrow Stream data."""
+
+        mimetype: ClassVar[str] = "application/vnd.apache.arrow.stream"
+
+        def loads(self, data: bytes) -> pa.Table:
+            try:
+                with pa.ipc.open_stream(data) as reader:
+                    return reader.read_all()
+            except pa.ArrowException as e:
+                raise SerializationError(
+                    message=f"Unable to deserialize as Arrow stream: {data!r}",
+                    errors=(e,),
+                )
+
+        def dumps(self, data: Any) -> bytes:
+            raise SerializationError(
+                message="Elasticsearch does not accept Arrow input data"
+            )
+
+
 DEFAULT_SERIALIZERS: Dict[str, Serializer] = {
     JsonSerializer.mimetype: JsonSerializer(),
     MapboxVectorTileSerializer.mimetype: MapboxVectorTileSerializer(),
@@ -122,6 +153,9 @@ DEFAULT_SERIALIZERS: Dict[str, Serializer] = {
     CompatibilityModeNdjsonSerializer.mimetype: CompatibilityModeNdjsonSerializer(),
 }
 
+if pa is not None:
+    DEFAULT_SERIALIZERS[PyArrowSerializer.mimetype] = PyArrowSerializer()
+
 # Alias for backwards compatibility
 JSONSerializer = JsonSerializer