VectorStore for GenAI integrations

dev-requirements.txt

@@ -15,10 +15,13 @@ twine
 build
 nox
 
-numpy
 pandas
 orjson
 
+# mmr for vectorstore
+numpy
+simsimd
+
 # Testing the 'search_mvt' API response
 mapbox-vector-tile
 # Python 3.7 gets an old version of mapbox-vector-tile, requiring an

elasticsearch/helpers/vectorstore/__init__.py

@@ -0,0 +1,62 @@
+#  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 elasticsearch.helpers.vectorstore._async.embedding_service import (
+    AsyncElasticsearchEmbeddings,
+    AsyncEmbeddingService,
+)
+from elasticsearch.helpers.vectorstore._async.strategies import (
+    AsyncBM25,
+    AsyncDenseVector,
+    AsyncDenseVectorScriptScore,
+    AsyncRetrievalStrategy,
+    AsyncSparseVector,
+)
+from elasticsearch.helpers.vectorstore._async.vectorstore import AsyncVectorStore
+from elasticsearch.helpers.vectorstore._sync.embedding_service import (
+    ElasticsearchEmbeddings,
+    EmbeddingService,
+)
+from elasticsearch.helpers.vectorstore._sync.strategies import (
+    BM25,
+    DenseVector,
+    DenseVectorScriptScore,
+    RetrievalStrategy,
+    SparseVector,
+)
+from elasticsearch.helpers.vectorstore._sync.vectorstore import VectorStore
+from elasticsearch.helpers.vectorstore._utils import DistanceMetric
+
+__all__ = [
+    "BM25",
+    "DenseVector",
+    "DenseVectorScriptScore",
+    "ElasticsearchEmbeddings",
+    "EmbeddingService",
+    "RetrievalStrategy",
+    "SparseVector",
+    "VectorStore",
+    "AsyncBM25",
+    "AsyncDenseVector",
+    "AsyncDenseVectorScriptScore",
+    "AsyncElasticsearchEmbeddings",
+    "AsyncEmbeddingService",
+    "AsyncRetrievalStrategy",
+    "AsyncSparseVector",
+    "AsyncVectorStore",
+    "DistanceMetric",
+]

elasticsearch/helpers/vectorstore/_async/__init__.py

@@ -0,0 +1,16 @@
+#  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.

elasticsearch/helpers/vectorstore/_async/_utils.py

@@ -0,0 +1,51 @@
+#  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 elasticsearch import (
+    AsyncElasticsearch,
+    BadRequestError,
+    ConflictError,
+    NotFoundError,
+)
+
+
+async def model_must_be_deployed(client: AsyncElasticsearch, model_id: str) -> None:
+    try:
+        dummy = {"x": "y"}
+        await client.ml.infer_trained_model(model_id=model_id, docs=[dummy])
+    except NotFoundError as err:
+        raise err
+    except ConflictError as err:
+        raise NotFoundError(
+            f"model '{model_id}' not found, please deploy it first",
+            meta=err.meta,
+            body=err.body,
+        ) from err
+    except BadRequestError:
+        # This error is expected because we do not know the expected document
+        # shape and just use a dummy doc above.
+        pass
+
+    return None
+
+
+async def model_is_deployed(es_client: AsyncElasticsearch, model_id: str) -> bool:
+    try:
+        await model_must_be_deployed(es_client, model_id)
+        return True
+    except NotFoundError:
+        return False

elasticsearch/helpers/vectorstore/_async/embedding_service.py

@@ -0,0 +1,97 @@
+#  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 List, Optional
+
+from elasticsearch import AsyncElasticsearch
+
+
+class AsyncEmbeddingService(ABC):
+    @abstractmethod
+    async def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Generate embeddings for a list of documents.
+
+        Args:
+            texts: A list of document strings to generate embeddings for.
+
+        Returns:
+            A list of embeddings, one for each document in the input.
+        """
+
+    @abstractmethod
+    async def embed_query(self, query: str) -> List[float]:
+        """Generate an embedding for a single query text.
+
+        Args:
+            text: The query text to generate an embedding for.
+
+        Returns:
+            The embedding for the input query text.
+        """
+
+
+class AsyncElasticsearchEmbeddings(AsyncEmbeddingService):
+    """Elasticsearch as a service for embedding model inference.
+
+    You need to have an embedding model downloaded and deployed in Elasticsearch:
+    - https://www.elastic.co/guide/en/elasticsearch/reference/current/infer-trained-model.html
+    - https://www.elastic.co/guide/en/machine-learning/current/ml-nlp-deploy-models.html
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        es_client: AsyncElasticsearch,
+        user_agent: str,
+        model_id: str,
+        input_field: str = "text_field",
+        num_dimensions: Optional[int] = None,
+    ):
+        """
+        Args:
+            agent_header: user agent header specific to the 3rd party integration.
+                Used for usage tracking in Elastic Cloud.
+            model_id: The model_id of the model deployed in the Elasticsearch cluster.
+            input_field: The name of the key for the input text field in the
+                document. Defaults to 'text_field'.
+            num_dimensions: The number of embedding dimensions. If None, then dimensions
+                will be infer from an example inference call.
+            es_client: Elasticsearch client connection. Alternatively specify the
+                Elasticsearch connection with the other es_* parameters.
+        """
+        # Add integration-specific usage header for tracking usage in Elastic Cloud.
+        # client.options preserces existing (non-user-agent) headers.
+        es_client = es_client.options(headers={"User-Agent": user_agent})
+
+        self.client = es_client.ml
+        self.model_id = model_id
+        self.input_field = input_field
+        self._num_dimensions = num_dimensions
+
+    async def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        result = await self._embedding_func(texts)
+        return result
+
+    async def embed_query(self, text: str) -> List[float]:
+        result = await self._embedding_func([text])
+        return result[0]
+
+    async def _embedding_func(self, texts: List[str]) -> List[List[float]]:
+        response = await self.client.infer_trained_model(
+            model_id=self.model_id, docs=[{self.input_field: text} for text in texts]
+        )
+        return [doc["predicted_value"] for doc in response["inference_results"]]