batch embedding call; infer num_dimensions

Author: maxjakob

elasticsearch/vectorstore/_async/strategies.py

@@ -21,7 +21,6 @@
 
 from elasticsearch import AsyncElasticsearch
 from elasticsearch.vectorstore._async._utils import model_must_be_deployed
-from elasticsearch.vectorstore._async.embedding_service import AsyncEmbeddingService
 
 
 class DistanceMetric(str, Enum):
@@ -63,7 +62,8 @@ async def create_index(
         self,
         client: AsyncElasticsearch,
         index_name: str,
-        metadata_mapping: Optional[Dict[str, str]],
+        num_dimensions: Optional[int] = None,
+        metadata_mapping: Optional[Dict[str, str]] = None,
     ) -> None:
         """
         Create the required index and do necessary preliminary work, like
@@ -76,21 +76,11 @@ async def create_index(
                 describe the schema of the metadata.
         """
 
-    async def embed_for_indexing(self, text: str) -> Dict[str, Any]:
+    def needs_inference(self) -> bool:
         """
-        If this strategy creates vector embeddings in Python (not in Elasticsearch),
-        this method is used to apply the inference.
-        The output is a dictionary with the vector field and the vector embedding.
-        It is merged in the ElasticserachStore with the rest of the document (text data,
-        metadata) before indexing.
-
-        Args:
-            text: Text input that can be used as input for inference.
-
-        Returns:
-            Dict: field and value pairs that extend the document to be indexed.
+        TODO
         """
-        return {}
+        return False
 
 
 # TODO test when repsective image is released
@@ -134,6 +124,7 @@ async def create_index(
         self,
         client: AsyncElasticsearch,
         index_name: str,
+        num_dimensions: int,
         metadata_mapping: Optional[Dict[str, str]],
     ) -> None:
         if self.model_id:
@@ -206,6 +197,7 @@ async def create_index(
         self,
         client: AsyncElasticsearch,
         index_name: str,
+        num_dimensions: int,
         metadata_mapping: Optional[Dict[str, str]],
     ) -> None:
         pipeline_name = f"{self.model_id}_sparse_embedding"
@@ -257,19 +249,11 @@ def __init__(
         knn_type: Literal["hnsw", "int8_hnsw", "flat", "int8_flat"] = "hnsw",
         vector_field: str = "vector_field",
         distance: DistanceMetric = DistanceMetric.COSINE,
-        embedding_service: Optional[AsyncEmbeddingService] = None,
         model_id: Optional[str] = None,
-        num_dimensions: Optional[int] = None,
         hybrid: bool = False,
         rrf: Union[bool, Dict[str, Any]] = True,
         text_field: Optional[str] = "text_field",
     ):
-        if embedding_service and model_id:
-            raise ValueError("either specify embedding_service or model_id, not both")
-        if model_id and not num_dimensions:
-            raise ValueError(
-                "if model_id is specified, num_dimensions must also be specified"
-            )
         if hybrid and not text_field:
             raise ValueError(
                 "to enable hybrid you have to specify a text_field (for BM25 matching)"
@@ -278,9 +262,7 @@ def __init__(
         self.knn_type = knn_type
         self.vector_field = vector_field
         self.distance = distance
-        self.embedding_service = embedding_service
         self.model_id = model_id
-        self.num_dimensions = num_dimensions
         self.hybrid = hybrid
         self.rrf = rrf
         self.text_field = text_field
@@ -302,10 +284,6 @@ async def es_query(
 
         if query_vector:
             knn["query_vector"] = query_vector
-        elif self.embedding_service:
-            knn["query_vector"] = await self.embedding_service.embed_query(
-                cast(str, query)
-            )
         else:
             # Inference in Elasticsearch. When initializing we make sure to always have
             # a model_id if don't have an embedding_service.
@@ -325,13 +303,9 @@ async def create_index(
         self,
         client: AsyncElasticsearch,
         index_name: str,
+        num_dimensions: int,
         metadata_mapping: Optional[Dict[str, str]],
     ) -> None:
-        if self.embedding_service and not self.num_dimensions:
-            self.num_dimensions = len(
-                await self.embedding_service.embed_query("get number of dimensions")
-            )
-
         if self.model_id:
             await model_must_be_deployed(client, self.model_id)
 
@@ -350,7 +324,7 @@ async def create_index(
             "properties": {
                 self.vector_field: {
                     "type": "dense_vector",
-                    "dims": self.num_dimensions,
+                    "dims": num_dimensions,
                     "index": True,
                     "similarity": similarityAlgo,
                 },
@@ -362,12 +336,6 @@ async def create_index(
         r = await client.indices.create(index=index_name, mappings=mappings)
         print(r)
 
-    async def embed_for_indexing(self, text: str) -> Dict[str, Any]:
-        if self.embedding_service:
-            vector = await self.embedding_service.embed_query(text)
-            return {self.vector_field: vector}
-        return {}
-
     def _hybrid(
         self, query: str, knn: Dict[str, Any], filter: List[Dict[str, Any]]
     ) -> Dict[str, Any]:
@@ -393,28 +361,27 @@ def _hybrid(
             },
         }
 
-        if isinstance(self.rrf, Dict[str, Any]):
+        if isinstance(self.rrf, Dict):
             query_body["rank"] = {"rrf": self.rrf}
         elif isinstance(self.rrf, bool) and self.rrf is True:
             query_body["rank"] = {"rrf": {}}
 
         return query_body
 
+    def needs_inference(self) -> bool:
+        return not self.model_id
+
 
 class DenseVectorScriptScore(RetrievalStrategy):
     """Exact nearest neighbors retrieval using the `script_score` query."""
 
     def __init__(
         self,
-        embedding_service: AsyncEmbeddingService,
         vector_field: str = "vector_field",
         distance: DistanceMetric = DistanceMetric.COSINE,
-        num_dimensions: Optional[int] = None,
     ) -> None:
         self.vector_field = vector_field
         self.distance = distance
-        self.embedding_service = embedding_service
-        self.num_dimensions = num_dimensions
 
     async def es_query(
         self,
@@ -424,6 +391,9 @@ async def es_query(
         filter: List[Dict[str, Any]] = [],
         query_vector: Optional[List[float]] = None,
     ) -> Dict[str, Any]:
+        if not query_vector:
+            raise ValueError("specify a query_vector")
+
         if self.distance is DistanceMetric.COSINE:
             similarityAlgo = (
                 f"cosineSimilarity(params.query_vector, '{self.vector_field}') + 1.0"
@@ -452,16 +422,6 @@ async def es_query(
         if filter:
             queryBool = {"bool": {"filter": filter}}
 
-        if not query_vector:
-            if not self.embedding_service:
-                raise ValueError(
-                    "if not embedding_service is given, you need to "
-                    "procive a query_vector"
-                )
-            if not query:
-                raise ValueError("either specify a query string or a query_vector")
-            query_vector = await self.embedding_service.embed_query(query)
-
         return {
             "query": {
                 "script_score": {
@@ -478,18 +438,14 @@ async def create_index(
         self,
         client: AsyncElasticsearch,
         index_name: str,
+        num_dimensions: int,
         metadata_mapping: Optional[Dict[str, str]],
     ) -> None:
-        if not self.num_dimensions:
-            self.num_dimensions = len(
-                await self.embedding_service.embed_query("get number of dimensions")
-            )
-
         mappings = {
             "properties": {
                 self.vector_field: {
                     "type": "dense_vector",
-                    "dims": self.num_dimensions,
+                    "dims": num_dimensions,
                     "index": False,
                 }
             }
@@ -499,10 +455,8 @@ async def create_index(
 
         await client.indices.create(index=index_name, mappings=mappings)
 
-        return None
-
-    async def embed_for_indexing(self, text: str) -> Dict[str, Any]:
-        return {self.vector_field: await self.embedding_service.embed_query(text)}
+    def needs_inference(self) -> bool:
+        return True
 
 
 class BM25(RetrievalStrategy):
@@ -545,6 +499,7 @@ async def create_index(
         self,
         client: AsyncElasticsearch,
         index_name: str,
+        num_dimensions: int,
         metadata_mapping: Optional[Dict[str, str]],
     ) -> None:
         similarity_name = "custom_bm25"

elasticsearch/vectorstore/_async/vectorestore.py

@@ -44,6 +44,8 @@ def __init__(
         user_agent: str,
         index_name: str,
         retrieval_strategy: RetrievalStrategy,
+        embedding_service: Optional[AsyncEmbeddingService] = None,
+        num_dimensions: Optional[int] = None,
         text_field: str = "text_field",
         vector_field: str = "vector_field",
         metadata_mapping: Optional[Dict[str, str]] = None,
@@ -61,7 +63,6 @@ def __init__(
             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})
@@ -74,6 +75,8 @@ def __init__(
         self.es_client = es_client
         self.index_name = index_name
         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_mapping = metadata_mapping
@@ -118,6 +121,9 @@ async def add_texts(
         if create_index_if_not_exists:
             await self._create_index_if_not_exists()
 
+        if self.embedding_service and not vectors:
+            vectors = await self.embedding_service.embed_documents(texts)
+
         for i, text in enumerate(texts):
             metadata = metadatas[i] if metadatas else {}
 
@@ -132,7 +138,6 @@ async def add_texts(
             if vectors:
                 request[self.vector_field] = vectors[i]
 
-            request.update(await self.retrieval_strategy.embed_for_indexing(text))
             requests.append(request)
 
         if len(requests) > 0:
@@ -240,6 +245,11 @@ async def search(
         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 = await self.embedding_service.embed_query(query)
+
         query_body = await self.retrieval_strategy.es_query(
             query=query,
             k=k,
@@ -267,9 +277,22 @@ async def _create_index_if_not_exists(self) -> None:
         if exists.meta.status == 200:
             logger.debug(f"Index {self.index_name} already exists. Skipping creation.")
         else:
+            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 = await self.embedding_service.embed_query(
+                        "get num dimensions"
+                    )
+                    self.num_dimensions = len(vector)
+
             await self.retrieval_strategy.create_index(
                 client=self.es_client,
                 index_name=self.index_name,
+                num_dimensions=self.num_dimensions,
                 metadata_mapping=self.metadata_mapping,
             )