elasticsearch9 9.0.0__py3-none-any.whl

elasticsearch9/helpers/vectorstore/_sync/vectorstore.py

@@ -0,0 +1,421 @@
+#  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 logging
+import uuid
+from typing import Any, Callable, Dict, List, Optional
+
+from elasticsearch import Elasticsearch
+from elasticsearch._version import __versionstr__ as lib_version
+from elasticsearch.helpers import BulkIndexError, bulk
+from elasticsearch.helpers.vectorstore import (
+    EmbeddingService,
+    RetrievalStrategy,
+)
+from elasticsearch.helpers.vectorstore._utils import maximal_marginal_relevance
+
+logger = logging.getLogger(__name__)
+
+
+class VectorStore:
+    """
+    VectorStore is a higher-level abstraction of indexing and search.
+    Users can pick from available retrieval strategies.
+
+    Documents have up to 3 fields:
+      - text_field: the text to be indexed and searched.
+      - metadata: additional information about the document, either schema-free
+        or defined by the supplied metadata_mappings.
+      - vector_field (usually not filled by the user): the embedding vector of the text.
+
+    Depending on the strategy, vector embeddings are
+      - created by the user beforehand
+      - created by this AsyncVectorStore class in Python
+      - created in-stack by inference pipelines.
+    """
+
+    def __init__(
+        self,
+        client: Elasticsearch,
+        *,
+        index: str,
+        retrieval_strategy: RetrievalStrategy,
+        embedding_service: Optional[EmbeddingService] = None,
+        num_dimensions: Optional[int] = None,
+        text_field: str = "text_field",
+        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.
+            Used for usage tracking in Elastic Cloud.
+        :param index: The name of the index to query.
+        :param retrieval_strategy: how to index and search the data. See the strategies
+            module for availble strategies.
+        :param text_field: Name of the field with the textual data.
+        :param vector_field: For strategies that perform embedding inference in Python,
+            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.
+        client = client.options(headers={"User-Agent": user_agent})
+
+        if hasattr(retrieval_strategy, "text_field"):
+            retrieval_strategy.text_field = text_field
+        if hasattr(retrieval_strategy, "vector_field"):
+            retrieval_strategy.vector_field = vector_field
+
+        self.client = client
+        self.index = index
+        self.retrieval_strategy = retrieval_strategy
+        self.embedding_service = embedding_service
+        self.num_dimensions = num_dimensions
+        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()
+
+    def add_texts(
+        self,
+        texts: List[str],
+        *,
+        metadatas: Optional[List[Dict[str, Any]]] = None,
+        vectors: Optional[List[List[float]]] = None,
+        ids: Optional[List[str]] = None,
+        refresh_indices: bool = True,
+        create_index_if_not_exists: bool = True,
+        bulk_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> List[str]:
+        """Add documents to the Elasticsearch index.
+
+        :param texts: List of text documents.
+        :param metadata: Optional list of document metadata. Must be of same length as
+            texts.
+        :param vectors: Optional list of embedding vectors. Must be of same length as
+            texts.
+        :param ids: Optional list of ID strings. Must be of same length as texts.
+        :param refresh_indices: Whether to refresh the index after deleting documents.
+            Defaults to True.
+        :param create_index_if_not_exists: Whether to create the index if it does not
+            exist. Defaults to True.
+        :param bulk_kwargs: Arguments to pass to the bulk function when indexing
+            (for example chunk_size).
+
+        :return: List of IDs of the created documents, either echoing the provided one
+            or returning newly created ones.
+        """
+        bulk_kwargs = bulk_kwargs or {}
+        ids = ids or [str(uuid.uuid4()) for _ in texts]
+        requests = []
+
+        if create_index_if_not_exists:
+            self._create_index_if_not_exists()
+
+        if self.embedding_service and not vectors:
+            vectors = self.embedding_service.embed_documents(texts)
+
+        for i, text in enumerate(texts):
+            metadata = metadatas[i] if metadatas else {}
+
+            request: Dict[str, Any] = {
+                "_op_type": "index",
+                "_index": self.index,
+                self.text_field: text,
+                "metadata": metadata,
+                "_id": ids[i],
+            }
+
+            if vectors:
+                request[self.vector_field] = vectors[i]
+
+            requests.append(request)
+
+        if len(requests) > 0:
+            try:
+                success, failed = bulk(
+                    self.client,
+                    requests,
+                    stats_only=True,
+                    refresh=refresh_indices,
+                    **bulk_kwargs,
+                )
+                logger.debug(f"added texts {ids} to index")
+                return ids
+            except BulkIndexError as e:
+                logger.error(f"Error adding texts: {e}")
+                firstError = e.errors[0].get("index", {}).get("error", {})
+                logger.error(f"First error reason: {firstError.get('reason')}")
+                raise e
+
+        else:
+            logger.debug("No texts to add to index")
+            return []
+
+    def delete(  # type: ignore[no-untyped-def]
+        self,
+        *,
+        ids: Optional[List[str]] = None,
+        query: Optional[Dict[str, Any]] = None,
+        refresh_indices: bool = True,
+        **delete_kwargs,
+    ) -> bool:
+        """Delete documents from the Elasticsearch index.
+
+        :param ids: List of IDs of documents to delete.
+        :param refresh_indices: Whether to refresh the index after deleting documents.
+            Defaults to True.
+
+        :return: True if deletion was successful.
+        """
+        if ids is not None and query is not None:
+            raise ValueError("one of ids or query must be specified")
+        elif ids is None and query is None:
+            raise ValueError("either specify ids or query")
+
+        try:
+            if ids:
+                body = [
+                    {"_op_type": "delete", "_index": self.index, "_id": _id}
+                    for _id in ids
+                ]
+                bulk(
+                    self.client,
+                    body,
+                    refresh=refresh_indices,
+                    ignore_status=404,
+                    **delete_kwargs,
+                )
+                logger.debug(f"Deleted {len(body)} texts from index")
+
+            else:
+                self.client.delete_by_query(
+                    index=self.index,
+                    query=query,
+                    refresh=refresh_indices,
+                    **delete_kwargs,
+                )
+
+        except BulkIndexError as e:
+            logger.error(f"Error deleting texts: {e}")
+            firstError = e.errors[0].get("index", {}).get("error", {})
+            logger.error(f"First error reason: {firstError.get('reason')}")
+            raise e
+
+        return True
+
+    def search(
+        self,
+        *,
+        query: Optional[str] = None,
+        query_vector: Optional[List[float]] = None,
+        k: int = 4,
+        num_candidates: int = 50,
+        fields: Optional[List[str]] = None,
+        filter: Optional[List[Dict[str, Any]]] = None,
+        custom_query: Optional[
+            Callable[[Dict[str, Any], Optional[str]], Dict[str, Any]]
+        ] = None,
+    ) -> List[Dict[str, Any]]:
+        """
+        :param query: Input query string.
+        :param query_vector: Input embedding vector. If given, input query string is
+            ignored.
+        :param k: Number of returned results.
+        :param num_candidates: Number of candidates to fetch from data nodes in knn.
+        :param fields: List of field names to return.
+        :param filter: Elasticsearch filters to apply.
+        :param custom_query: Function to modify the Elasticsearch query body before it is
+            sent to Elasticsearch.
+
+        :return: List of document hits. Includes _index, _id, _score and _source.
+        """
+        if fields is None:
+            fields = []
+        if "metadata" not in fields:
+            fields.append("metadata")
+        if self.text_field not in fields:
+            fields.append(self.text_field)
+
+        if self.embedding_service and not query_vector:
+            if not query:
+                raise ValueError("specify a query or a query_vector to search")
+            query_vector = self.embedding_service.embed_query(query)
+
+        query_body = self.retrieval_strategy.es_query(
+            query=query,
+            query_vector=query_vector,
+            text_field=self.text_field,
+            vector_field=self.vector_field,
+            k=k,
+            num_candidates=num_candidates,
+            filter=filter or [],
+        )
+
+        if custom_query is not None:
+            query_body = custom_query(query_body, query)
+            logger.debug(f"Calling custom_query, Query body now: {query_body}")
+
+        response = self.client.search(
+            index=self.index,
+            **query_body,
+            size=k,
+            source=True,
+            source_includes=fields,
+        )
+        hits: List[Dict[str, Any]] = response["hits"]["hits"]
+
+        return hits
+
+    def _create_index_if_not_exists(self) -> None:
+        exists = self.client.indices.exists(index=self.index)
+        if exists.meta.status == 200:
+            logger.debug(f"Index {self.index} already exists. Skipping creation.")
+            return
+
+        if self.retrieval_strategy.needs_inference():
+            if not self.num_dimensions and not self.embedding_service:
+                raise ValueError(
+                    "retrieval strategy requires embeddings; either embedding_service "
+                    "or num_dimensions need to be specified"
+                )
+            if not self.num_dimensions and self.embedding_service:
+                vector = self.embedding_service.embed_query("get num dimensions")
+                self.num_dimensions = len(vector)
+
+        mappings, settings = self.retrieval_strategy.es_mappings_settings(
+            text_field=self.text_field,
+            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():
+                if key in metadata:
+                    raise ValueError(f"metadata key {key} already exists in mappings")
+
+            metadata = dict(**metadata["properties"], **self.metadata_mappings)
+            mappings["properties"]["metadata"] = {"properties": metadata}
+
+        self.retrieval_strategy.before_index_creation(
+            client=self.client,
+            text_field=self.text_field,
+            vector_field=self.vector_field,
+        )
+        self.client.indices.create(
+            index=self.index, mappings=mappings, settings=settings
+        )
+
+    def max_marginal_relevance_search(
+        self,
+        *,
+        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,
+        lambda_mult: float = 0.5,
+        fields: Optional[List[str]] = None,
+        custom_query: Optional[
+            Callable[[Dict[str, Any], Optional[str]], Dict[str, Any]]
+        ] = None,
+    ) -> List[Dict[str, Any]]:
+        """Return docs selected using the maximal marginal relevance.
+
+        Maximal marginal relevance optimizes for similarity to query AND diversity
+            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
+            of diversity among the results with 0 corresponding
+            to maximum diversity and 1 to minimum diversity.
+            Defaults to 0.5.
+        :param fields: Other fields to get from elasticsearch source. These fields
+            will be added to the document metadata.
+
+        :return: A list of Documents selected by maximal marginal relevance.
+        """
+        remove_vector_query_field_from_metadata = True
+        if fields is None:
+            fields = [vector_field]
+        elif vector_field not in fields:
+            fields.append(vector_field)
+        else:
+            remove_vector_query_field_from_metadata = False
+
+        # Embed the 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_vector,
+            k=num_candidates,
+            fields=fields,
+            custom_query=custom_query,
+        )
+
+        # Get the embeddings for the fetched documents
+        got_embeddings = [hit["_source"][vector_field] for hit in got_hits]
+
+        # Select documents using maximal marginal relevance
+        selected_indices = maximal_marginal_relevance(
+            query_vector, got_embeddings, lambda_mult=lambda_mult, k=k
+        )
+        selected_hits = [got_hits[i] for i in selected_indices]
+
+        if remove_vector_query_field_from_metadata:
+            for hit in selected_hits:
+                del hit["_source"][vector_field]
+
+        return selected_hits

elasticsearch9/helpers/vectorstore/_utils.py

@@ -0,0 +1,116 @@
+#  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 enum import Enum
+from typing import TYPE_CHECKING, List, Union
+
+if TYPE_CHECKING:
+    import numpy as np
+    import numpy.typing as npt
+
+Matrix = Union[
+    List[List[float]], List["npt.NDArray[np.float64]"], "npt.NDArray[np.float64]"
+]
+
+
+class DistanceMetric(str, Enum):
+    """Enumerator of all Elasticsearch dense vector distance metrics."""
+
+    COSINE = "COSINE"
+    DOT_PRODUCT = "DOT_PRODUCT"
+    EUCLIDEAN_DISTANCE = "EUCLIDEAN_DISTANCE"
+    MAX_INNER_PRODUCT = "MAX_INNER_PRODUCT"
+
+
+def maximal_marginal_relevance(
+    query_embedding: List[float],
+    embedding_list: List[List[float]],
+    lambda_mult: float = 0.5,
+    k: int = 4,
+) -> List[int]:
+    """Calculate maximal marginal relevance."""
+
+    try:
+        import numpy as np
+    except ModuleNotFoundError as e:
+        _raise_missing_mmr_deps_error(e)
+
+    query_embedding_arr = np.array(query_embedding)
+
+    if min(k, len(embedding_list)) <= 0:
+        return []
+    if query_embedding_arr.ndim == 1:
+        query_embedding_arr = np.expand_dims(query_embedding_arr, axis=0)
+    similarity_to_query = _cosine_similarity(query_embedding_arr, embedding_list)[0]
+    most_similar = int(np.argmax(similarity_to_query))
+    idxs = [most_similar]
+    selected = np.array([embedding_list[most_similar]])
+    while len(idxs) < min(k, len(embedding_list)):
+        best_score = -np.inf
+        idx_to_add = -1
+        similarity_to_selected = _cosine_similarity(embedding_list, selected)
+        for i, query_score in enumerate(similarity_to_query):
+            if i in idxs:
+                continue
+            redundant_score = max(similarity_to_selected[i])
+            equation_score = (
+                lambda_mult * query_score - (1 - lambda_mult) * redundant_score
+            )
+            if equation_score > best_score:
+                best_score = equation_score
+                idx_to_add = i
+        idxs.append(idx_to_add)
+        selected = np.append(selected, [embedding_list[idx_to_add]], axis=0)
+    return idxs
+
+
+def _cosine_similarity(X: Matrix, Y: Matrix) -> "npt.NDArray[np.float64]":
+    """Row-wise cosine similarity between two equal-width matrices."""
+
+    try:
+        import numpy as np
+        import simsimd as simd
+    except ModuleNotFoundError as e:
+        _raise_missing_mmr_deps_error(e)
+
+    if len(X) == 0 or len(Y) == 0:
+        return np.array([])
+
+    X = np.array(X)
+    Y = np.array(Y)
+    if X.shape[1] != Y.shape[1]:
+        raise ValueError(
+            f"Number of columns in X and Y must be the same. X has shape {X.shape} "
+            f"and Y has shape {Y.shape}."
+        )
+
+    X = np.array(X, dtype=np.float32)
+    Y = np.array(Y, dtype=np.float32)
+    Z = 1 - np.array(simd.cdist(X, Y, metric="cosine"))
+    if isinstance(Z, float):
+        return np.array([Z])
+    return np.array(Z)
+
+
+def _raise_missing_mmr_deps_error(parent_error: ModuleNotFoundError) -> None:
+    import sys
+
+    raise ModuleNotFoundError(
+        f"Failed to compute maximal marginal relevance because the required "
+        f"module '{parent_error.name}' is missing. You can install it by running: "
+        f"'{sys.executable} -m pip install elasticsearch[vectorstore_mmr]'"
+    ) from parent_error