elasticsearch9 9.0.0__py3-none-any.whl

elasticsearch9/helpers/vectorstore/_sync/strategies.py

@@ -0,0 +1,487 @@
+#  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 abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Tuple, Union, cast
+
+from elasticsearch import Elasticsearch
+from elasticsearch.helpers.vectorstore._sync._utils import model_must_be_deployed
+from elasticsearch.helpers.vectorstore._utils import DistanceMetric
+
+
+class RetrievalStrategy(ABC):
+    @abstractmethod
+    def es_query(
+        self,
+        *,
+        query: Optional[str],
+        query_vector: Optional[List[float]],
+        text_field: str,
+        vector_field: str,
+        k: int,
+        num_candidates: int,
+        filter: List[Dict[str, Any]] = [],
+    ) -> Dict[str, Any]:
+        """
+        Returns the Elasticsearch query body for the given parameters.
+        The store will execute the query.
+
+        :param query: The text query. Can be None if query_vector is given.
+        :param k: The total number of results to retrieve.
+        :param num_candidates: The number of results to fetch initially in knn search.
+        :param filter: List of filter clauses to apply to the query.
+        :param query_vector: The query vector. Can be None if a query string is given.
+
+        :return: The Elasticsearch query body.
+        """
+
+    @abstractmethod
+    def es_mappings_settings(
+        self,
+        *,
+        text_field: str,
+        vector_field: str,
+        num_dimensions: Optional[int],
+    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        """
+        Create the required index and do necessary preliminary work, like
+        creating inference pipelines or checking if a required model was deployed.
+
+        :param client: Elasticsearch client connection.
+        :param text_field: The field containing the text data in the index.
+        :param vector_field: The field containing the vector representations in the index.
+        :param num_dimensions: If vectors are indexed, how many dimensions do they have.
+
+        :return: Dictionary with field and field type pairs that describe the schema.
+        """
+
+    def before_index_creation(
+        self, *, client: Elasticsearch, text_field: str, vector_field: str
+    ) -> None:
+        """
+        Executes before the index is created. Used for setting up
+        any required Elasticsearch resources like a pipeline.
+        Defaults to a no-op.
+
+        :param client: The Elasticsearch client.
+        :param text_field: The field containing the text data in the index.
+        :param vector_field: The field containing the vector representations in the index.
+        """
+        pass
+
+    def needs_inference(self) -> bool:
+        """
+        Some retrieval strategies index embedding vectors and allow search by embedding
+        vector, for example the `DenseVectorStrategy` strategy. Mapping a user input query
+        string to an embedding vector is called inference. Inference can be applied
+        in Elasticsearch (using a `model_id`) or outside of Elasticsearch (using an
+        `EmbeddingService` defined on the `VectorStore`). In the latter case,
+        this method has to return True.
+        """
+        return False
+
+
+class SparseVectorStrategy(RetrievalStrategy):
+    """Sparse retrieval strategy using the `sparse_vector` processor."""
+
+    def __init__(self, model_id: str = ".elser_model_2"):
+        self.model_id = model_id
+        self._tokens_field = "tokens"
+        self._pipeline_name = f"{self.model_id}_sparse_embedding"
+
+    def es_query(
+        self,
+        *,
+        query: Optional[str],
+        query_vector: Optional[List[float]],
+        text_field: str,
+        vector_field: str,
+        k: int,
+        num_candidates: int,
+        filter: List[Dict[str, Any]] = [],
+    ) -> Dict[str, Any]:
+        if query_vector:
+            raise ValueError(
+                "Cannot do sparse retrieval with a query_vector. "
+                "Inference is currently always applied in Elasticsearch."
+            )
+        if query is None:
+            raise ValueError("please specify a query string")
+
+        return {
+            "query": {
+                "bool": {
+                    "must": [
+                        {
+                            "sparse_vector": {
+                                "field": f"{vector_field}.{self._tokens_field}",
+                                "inference_id": self.model_id,
+                                "query": query,
+                            }
+                        }
+                    ],
+                    "filter": filter,
+                }
+            }
+        }
+
+    def es_mappings_settings(
+        self,
+        *,
+        text_field: str,
+        vector_field: str,
+        num_dimensions: Optional[int],
+    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        mappings: Dict[str, Any] = {
+            "properties": {
+                vector_field: {
+                    "properties": {self._tokens_field: {"type": "sparse_vector"}}
+                }
+            }
+        }
+        settings = {"default_pipeline": self._pipeline_name}
+
+        return mappings, settings
+
+    def before_index_creation(
+        self, *, client: Elasticsearch, text_field: str, vector_field: str
+    ) -> None:
+        if self.model_id:
+            model_must_be_deployed(client, self.model_id)
+
+            # Create a pipeline for the model
+            client.ingest.put_pipeline(
+                id=self._pipeline_name,
+                description="Embedding pipeline for Python VectorStore",
+                processors=[
+                    {
+                        "inference": {
+                            "model_id": self.model_id,
+                            "input_output": [
+                                {
+                                    "input_field": text_field,
+                                    "output_field": f"{vector_field}.{self._tokens_field}",
+                                },
+                            ],
+                        }
+                    }
+                ],
+            )
+
+
+class DenseVectorStrategy(RetrievalStrategy):
+    """K-nearest-neighbors retrieval."""
+
+    def __init__(
+        self,
+        *,
+        distance: DistanceMetric = DistanceMetric.COSINE,
+        model_id: Optional[str] = None,
+        hybrid: bool = False,
+        rrf: Union[bool, Dict[str, Any]] = True,
+        text_field: Optional[str] = "text_field",
+    ):
+        if hybrid and not text_field:
+            raise ValueError(
+                "to enable hybrid you have to specify a text_field (for BM25Strategy matching)"
+            )
+
+        self.distance = distance
+        self.model_id = model_id
+        self.hybrid = hybrid
+        self.rrf = rrf
+        self.text_field = text_field
+
+    def es_query(
+        self,
+        *,
+        query: Optional[str],
+        query_vector: Optional[List[float]],
+        text_field: str,
+        vector_field: str,
+        k: int,
+        num_candidates: int,
+        filter: List[Dict[str, Any]] = [],
+    ) -> Dict[str, Any]:
+        knn = {
+            "filter": filter,
+            "field": vector_field,
+            "k": k,
+            "num_candidates": num_candidates,
+        }
+
+        if query_vector is not None:
+            knn["query_vector"] = query_vector
+        else:
+            # Inference in Elasticsearch. When initializing we make sure to always have
+            # a model_id if don't have an embedding_service.
+            knn["query_vector_builder"] = {
+                "text_embedding": {
+                    "model_id": self.model_id,
+                    "model_text": query,
+                }
+            }
+
+        if self.hybrid:
+            return self._hybrid(query=cast(str, query), knn=knn, filter=filter)
+
+        return {"knn": knn}
+
+    def es_mappings_settings(
+        self,
+        *,
+        text_field: str,
+        vector_field: str,
+        num_dimensions: Optional[int],
+    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        if self.distance is DistanceMetric.COSINE:
+            similarity = "cosine"
+        elif self.distance is DistanceMetric.EUCLIDEAN_DISTANCE:
+            similarity = "l2_norm"
+        elif self.distance is DistanceMetric.DOT_PRODUCT:
+            similarity = "dot_product"
+        elif self.distance is DistanceMetric.MAX_INNER_PRODUCT:
+            similarity = "max_inner_product"
+        else:
+            raise ValueError(f"Similarity {self.distance} not supported.")
+
+        mappings: Dict[str, Any] = {
+            "properties": {
+                vector_field: {
+                    "type": "dense_vector",
+                    "dims": num_dimensions,
+                    "index": True,
+                    "similarity": similarity,
+                },
+            }
+        }
+
+        return mappings, {}
+
+    def before_index_creation(
+        self, *, client: Elasticsearch, text_field: str, vector_field: str
+    ) -> None:
+        if self.model_id:
+            model_must_be_deployed(client, self.model_id)
+
+    def _hybrid(
+        self, query: str, knn: Dict[str, Any], filter: List[Dict[str, Any]]
+    ) -> Dict[str, Any]:
+        # Add a query to the knn query.
+        # RRF is used to even the score from the knn query and text query
+        # RRF has two optional parameters: {'rank_constant':int, 'rank_window_size':int}
+        # https://www.elastic.co/guide/en/elasticsearch/reference/current/rrf.html
+        standard_query = {
+            "query": {
+                "bool": {
+                    "must": [
+                        {
+                            "match": {
+                                self.text_field: {
+                                    "query": query,
+                                }
+                            }
+                        }
+                    ],
+                    "filter": filter,
+                }
+            }
+        }
+
+        if self.rrf is False:
+            query_body = {
+                "knn": knn,
+                **standard_query,
+            }
+        else:
+            rrf_options = {}
+            if isinstance(self.rrf, Dict):
+                if "rank_constant" in self.rrf:
+                    rrf_options["rank_constant"] = self.rrf["rank_constant"]
+                if "window_size" in self.rrf:
+                    # 'window_size' was renamed to 'rank_window_size', but we support
+                    # the older name for backwards compatibility
+                    rrf_options["rank_window_size"] = self.rrf["window_size"]
+                if "rank_window_size" in self.rrf:
+                    rrf_options["rank_window_size"] = self.rrf["rank_window_size"]
+            query_body = {
+                "retriever": {
+                    "rrf": {
+                        "retrievers": [
+                            {"standard": standard_query},
+                            {"knn": knn},
+                        ],
+                        **rrf_options,
+                    },
+                },
+            }
+        return query_body
+
+    def needs_inference(self) -> bool:
+        return not self.model_id
+
+
+class DenseVectorScriptScoreStrategy(RetrievalStrategy):
+    """Exact nearest neighbors retrieval using the `script_score` query."""
+
+    def __init__(self, distance: DistanceMetric = DistanceMetric.COSINE) -> None:
+        self.distance = distance
+
+    def es_query(
+        self,
+        *,
+        query: Optional[str],
+        query_vector: Optional[List[float]],
+        text_field: str,
+        vector_field: str,
+        k: int,
+        num_candidates: int,
+        filter: List[Dict[str, Any]] = [],
+    ) -> Dict[str, Any]:
+        if not query_vector:
+            raise ValueError("specify a query_vector")
+
+        if self.distance is DistanceMetric.COSINE:
+            similarity_algo = (
+                f"cosineSimilarity(params.query_vector, '{vector_field}') + 1.0"
+            )
+        elif self.distance is DistanceMetric.EUCLIDEAN_DISTANCE:
+            similarity_algo = f"1 / (1 + l2norm(params.query_vector, '{vector_field}'))"
+        elif self.distance is DistanceMetric.DOT_PRODUCT:
+            similarity_algo = f"""
+            double value = dotProduct(params.query_vector, '{vector_field}');
+            return sigmoid(1, Math.E, -value);
+            """
+        elif self.distance is DistanceMetric.MAX_INNER_PRODUCT:
+            similarity_algo = f"""
+            double value = dotProduct(params.query_vector, '{vector_field}');
+            if (dotProduct < 0) {{
+                return 1 / (1 + -1 * dotProduct);
+            }}
+            return dotProduct + 1;
+            """
+        else:
+            raise ValueError(f"Similarity {self.distance} not supported.")
+
+        query_bool: Dict[str, Any] = {"match_all": {}}
+        if filter:
+            query_bool = {"bool": {"filter": filter}}
+
+        return {
+            "query": {
+                "script_score": {
+                    "query": query_bool,
+                    "script": {
+                        "source": similarity_algo,
+                        "params": {"query_vector": query_vector},
+                    },
+                },
+            }
+        }
+
+    def es_mappings_settings(
+        self,
+        *,
+        text_field: str,
+        vector_field: str,
+        num_dimensions: Optional[int],
+    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        mappings = {
+            "properties": {
+                vector_field: {
+                    "type": "dense_vector",
+                    "dims": num_dimensions,
+                    "index": False,
+                }
+            }
+        }
+
+        return mappings, {}
+
+    def needs_inference(self) -> bool:
+        return True
+
+
+class BM25Strategy(RetrievalStrategy):
+    def __init__(
+        self,
+        k1: Optional[float] = None,
+        b: Optional[float] = None,
+    ):
+        self.k1 = k1
+        self.b = b
+
+    def es_query(
+        self,
+        *,
+        query: Optional[str],
+        query_vector: Optional[List[float]],
+        text_field: str,
+        vector_field: str,
+        k: int,
+        num_candidates: int,
+        filter: List[Dict[str, Any]] = [],
+    ) -> Dict[str, Any]:
+        return {
+            "query": {
+                "bool": {
+                    "must": [
+                        {
+                            "match": {
+                                text_field: {
+                                    "query": query,
+                                }
+                            },
+                        },
+                    ],
+                    "filter": filter,
+                },
+            },
+        }
+
+    def es_mappings_settings(
+        self,
+        *,
+        text_field: str,
+        vector_field: str,
+        num_dimensions: Optional[int],
+    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        similarity_name = "custom_bm25"
+
+        mappings: Dict[str, Any] = {
+            "properties": {
+                text_field: {
+                    "type": "text",
+                    "similarity": similarity_name,
+                },
+            },
+        }
+
+        bm25: Dict[str, Any] = {
+            "type": "BM25",
+        }
+        if self.k1 is not None:
+            bm25["k1"] = self.k1
+        if self.b is not None:
+            bm25["b"] = self.b
+        settings = {
+            "similarity": {
+                similarity_name: bm25,
+            }
+        }
+
+        return mappings, settings