Add support for LangChain embeddings in Object API (#280)

This adds support for LangChain embeddings in Object API.

It also adds support for passing environment variables to the ingestion execution.

Author: NikolaosPapailiou

apis/python/examples/object_api/text_search_documents.ipynb

@@ -22,7 +22,7 @@
     "import tiledb\n",
     "from tiledb.vector_search.object_api import object_index\n",
     "from tiledb.vector_search.object_readers import DirectoryTextReader\n",
-    "from tiledb.vector_search.embeddings import SentenceTransformersEmbedding\n",
+    "from tiledb.vector_search.embeddings import SentenceTransformersEmbedding, LangChainEmbedding\n",
     "\n",
     "dataset = \"documents\"\n",
     "base_uri = f\"/tmp/{dataset}_demo\"\n",
@@ -87,6 +87,20 @@
     "        text_splitter_kwargs={\"chunk_size\":1000}\n",
     "    )\n",
     "embedding = SentenceTransformersEmbedding(model_name_or_path='BAAI/bge-small-en-v1.5', dimensions=384)\n",
+    "# embedding = LangChainEmbedding(\n",
+    "#     dimensions=384, \n",
+    "#     embedding_class=\"HuggingFaceEmbeddings\", \n",
+    "#     embedding_kwargs={\n",
+    "#         \"model_name\": 'BAAI/bge-small-en-v1.5', \n",
+    "#         }\n",
+    "# )\n",
+    "# embedding = LangChainEmbedding(\n",
+    "#     dimensions=1536, \n",
+    "#     embedding_class=\"OpenAIEmbeddings\", \n",
+    "#     embedding_kwargs={\n",
+    "#         \"model\": 'text-embedding-ada-002', \n",
+    "#         }\n",
+    "# )\n",
     "index = object_index.create(\n",
     "    uri=index_uri,\n",
     "    index_type=\"IVF_FLAT\",\n",

apis/python/src/tiledb/vector_search/embeddings/__init__.py

@@ -1,4 +1,5 @@
 from .image_resnetv2_embedding import ImageResNetV2Embedding
+from .langchain_embedding import LangChainEmbedding
 from .object_embedding import ObjectEmbedding
 from .random_embedding import RandomEmbedding
 from .sentence_transformers_embedding import SentenceTransformersEmbedding
@@ -10,4 +11,5 @@
     "ImageResNetV2Embedding",
     "RandomEmbedding",
     "SentenceTransformersEmbedding",
+    "LangChainEmbedding",
 ]

apis/python/src/tiledb/vector_search/embeddings/langchain_embedding.py

@@ -0,0 +1,47 @@
+from typing import Dict, Optional, OrderedDict
+
+import numpy as np
+
+# from tiledb.vector_search.embeddings import ObjectEmbedding
+
+
+# class LangChainEmbedding(ObjectEmbedding):
+class LangChainEmbedding():
+    """
+    Embedding functions from `langchain.embeddings` package.
+    """
+
+    def __init__(
+        self,
+        dimensions: int,
+        embedding_class: str = "OpenAIEmbeddings",
+        embedding_kwargs: Optional[Dict] = None,
+    ):
+        self.dim_num = dimensions
+        self.embedding_class = embedding_class
+        self.embedding_kwargs = embedding_kwargs
+
+    def init_kwargs(self) -> Dict:
+        return {
+            "dimensions": self.dim_num,
+            "embedding_class": self.embedding_class,
+            "embedding_kwargs": self.embedding_kwargs,
+        }
+
+    def dimensions(self) -> int:
+        return self.dim_num
+
+    def vector_type(self) -> np.dtype:
+        return np.float32
+
+    def load(self) -> None:
+        import importlib
+
+        embeddings_module = importlib.import_module("langchain.embeddings")
+        embedding_class_ = getattr(embeddings_module, self.embedding_class)
+        self.embedding = embedding_class_(**self.embedding_kwargs)
+
+    def embed(self, objects: OrderedDict, metadata: OrderedDict) -> np.ndarray:
+        return np.array(
+            self.embedding.embed_documents(objects["text"]), dtype=np.float32
+        )

apis/python/src/tiledb/vector_search/object_api/embeddings_ingestion.py

@@ -25,6 +25,7 @@ def ingest_embeddings_with_driver(
     vector_indexing_mode: Mode = Mode.LOCAL,
     config: Optional[Mapping[str, Any]] = None,
     namespace: Optional[str] = None,
+    environment_variables: Dict = {},
     **kwargs,
 ):
     def ingest_embeddings(
@@ -45,6 +46,7 @@ def ingest_embeddings(
         vector_indexing_mode: Mode = Mode.LOCAL,
         config: Optional[Mapping[str, Any]] = None,
         namespace: Optional[str] = None,
+        environment_variables: Dict = {},
         **kwargs,
     ):
         import tiledb
@@ -133,6 +135,7 @@ def compute_embeddings_udf(
             trace_id: Optional[str] = None,
             config: Optional[Mapping[str, Any]] = None,
             extra_worker_modules: Optional[List[str]] = None,
+            environment_variables: Dict = {},
         ):
             def install_extra_driver_modules():
                 if extra_worker_modules is not None:
@@ -150,6 +153,8 @@ def install_extra_driver_modules():
 
             install_extra_driver_modules()
 
+            import os
+
             import numpy as np
 
             import tiledb
@@ -187,6 +192,8 @@ def instantiate_object(code, class_name, **kwargs):
                 class_name=object_embedding_class_name,
                 **object_embedding_kwargs,
             )
+            for var, val in environment_variables.items():
+                os.environ[var] = val
             with tiledb.scope_ctx(ctx_or_config=config):
                 logger.debug("Loading model...")
                 object_embedding.load()
@@ -263,6 +270,7 @@ def create_dag(
             trace_id: Optional[str] = None,
             embeddings_generation_mode: Mode = Mode.LOCAL,
             config: Optional[Mapping[str, Any]] = None,
+            environment_variables: Dict = {},
             namespace: Optional[str] = None,
         ) -> dag.DAG:
             if embeddings_generation_mode == Mode.BATCH:
@@ -327,6 +335,7 @@ def create_dag(
                     verbose=verbose,
                     trace_id=trace_id,
                     config=config,
+                    environment_variables=environment_variables,
                     extra_worker_modules=extra_udf_worker_modules,
                     name="generate-embeddings-" + str(task_id),
                     resources=worker_resources,
@@ -340,6 +349,10 @@ def create_dag(
         # End internal function definitions
         # --------------------------------------------------------------------
 
+        import os
+
+        for var, val in environment_variables.items():
+            os.environ[var] = val
         with tiledb.scope_ctx(ctx_or_config=config):
             logger = setup(config, verbose)
             logger.debug("Generating embeddings")
@@ -348,7 +361,11 @@ def create_dag(
 
             from tiledb.vector_search.object_api import object_index
 
-            ob_index = object_index.ObjectIndex(object_index_uri, config=config)
+            ob_index = object_index.ObjectIndex(
+                object_index_uri,
+                config=config,
+                environment_variables=environment_variables,
+            )
             partitions = ob_index.object_reader.get_partitions(**kwargs)
             object_partitions = len(partitions)
             object_partitions_per_worker = 1
@@ -394,6 +411,7 @@ def create_dag(
                 trace_id=trace_id,
                 embeddings_generation_mode=embeddings_generation_mode,
                 config=config,
+                environment_variables=environment_variables,
                 namespace=namespace,
             )
             logger.debug("Submitting ingestion graph")
@@ -465,6 +483,7 @@ def submit_local(d, func, *args, **kwargs):
         embeddings_generation_mode=embeddings_generation_mode,
         vector_indexing_mode=vector_indexing_mode,
         config=config,
+        environment_variables=environment_variables,
         **kwargs,
         name="ingest-embeddings-driver",
         resources={"cpu": "1", "memory": "4Gi"}

apis/python/src/tiledb/vector_search/object_api/object_index.py

@@ -28,8 +28,13 @@ def __init__(
         timestamp=None,
         load_embedding: bool = True,
         load_metadata_in_memory: bool = True,
+        environment_variables: Dict = {},
         **kwargs,
     ):
+        import os
+
+        for var, val in environment_variables.items():
+            os.environ[var] = val
         with tiledb.scope_ctx(ctx_or_config=config):
             self.uri = uri
             self.config = config
@@ -235,6 +240,7 @@ def update_index(
         vector_indexing_mode: Mode = Mode.LOCAL,
         config: Optional[Mapping[str, Any]] = None,
         namespace: Optional[str] = None,
+        environment_variables: Dict = {},
         **kwargs,
     ):
         embeddings_array_name = storage_formats[self.index.storage_version][
@@ -268,6 +274,7 @@ def update_index(
             embeddings_generation_mode=embeddings_generation_mode,
             vector_indexing_mode=vector_indexing_mode,
             config=config,
+            environment_variables=environment_variables,
             **kwargs,
         )