ref: adds multi threading to the AI/ML embeddings component (#2959)

* Use http client for requests and split texts naively

* update models list

* prints

* multithread requests to aiml embeddings

* remove comment

* [autofix.ci] apply automated fixes

* style(AIMLEmbeddingsImpl.py): improve code formatting and add type hinting for better code clarit

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>

src/backend/base/langflow/base/models/aiml_constants.py

@@ -1,4 +1,4 @@
-CHAT_MODELS = [
+AIML_CHAT_MODELS = [
     "zero-one-ai/Yi-34B-Chat",
     "allenai/OLMo-7B-Instruct",
     "allenai/OLMo-7B-Twin-2T",
@@ -17,6 +17,9 @@ CHAT_MODELS = [
     "codellama/CodeLlama-34b-Instruct-hf",
     "codellama/CodeLlama-70b-Instruct-hf",
     "codellama/CodeLlama-7b-Instruct-hf",
+    "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo",
+    "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
+    "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
     "meta-llama/Llama-2-70b-chat-hf",
     "meta-llama/Llama-2-13b-chat-hf",
     "meta-llama/Llama-2-7b-chat-hf",
@@ -70,7 +73,7 @@ CHAT_MODELS = [
     "claude-3-haiku-20240307",
 ]
 
-EMBEDDING_MODELS = [
+AIML_EMBEDDING_MODELS = [
     "text-embedding-3-small",
     "text-embedding-3-large",
     "text-embedding-ada-002",

src/backend/base/langflow/base/models/openai_constants.py

@@ -1,4 +1,4 @@
-MODEL_NAMES = [
+OPENAI_MODEL_NAMES = [
     "gpt-4o-mini",
     "gpt-4o",
     "gpt-4-turbo",
@@ -7,3 +7,9 @@ MODEL_NAMES = [
     "gpt-3.5-turbo",
     "gpt-3.5-turbo-0125",
 ]
+
+OPENAI_EMBEDDING_MODEL_NAMES = [
+    "text-embedding-3-small",
+    "text-embedding-3-large",
+    "text-embedding-ada-002",
+]

src/backend/base/langflow/components/embeddings/AIMLEmbeddings.py

@@ -1,5 +1,5 @@
 from langflow.base.embeddings.model import LCEmbeddingsModel
-from langflow.base.models.aiml_constants import EMBEDDING_MODELS
+from langflow.base.models.aiml_constants import AIML_EMBEDDING_MODELS
 from langflow.components.embeddings.util.AIMLEmbeddingsImpl import AIMLEmbeddingsImpl
 from langflow.field_typing import Embeddings
 from langflow.inputs.inputs import DropdownInput
@@ -16,7 +16,7 @@ class AIMLEmbeddingsComponent(LCEmbeddingsModel):
         DropdownInput(
             name="model_name",
             display_name="Model Name",
-            options=EMBEDDING_MODELS,
+            options=AIML_EMBEDDING_MODELS,
             required=True,
         ),
         SecretStrInput(

src/backend/base/langflow/components/embeddings/OpenAIEmbeddings.py

@@ -1,6 +1,7 @@
 from langchain_openai.embeddings.base import OpenAIEmbeddings
 
 from langflow.base.embeddings.model import LCEmbeddingsModel
+from langflow.base.models.openai_constants import OPENAI_EMBEDDING_MODEL_NAMES
 from langflow.field_typing import Embeddings
 from langflow.io import BoolInput, DictInput, DropdownInput, FloatInput, IntInput, MessageTextInput, SecretStrInput
 
@@ -33,11 +34,7 @@ class OpenAIEmbeddingsComponent(LCEmbeddingsModel):
             name="model",
             display_name="Model",
             advanced=False,
-            options=[
+            options=OPENAI_EMBEDDING_MODEL_NAMES,
-                "text-embedding-3-small",
-                "text-embedding-3-large",
-                "text-embedding-ada-002",
-            ],
             value="text-embedding-3-small",
         ),
         DictInput(name="model_kwargs", display_name="Model Kwargs", advanced=True),

src/backend/base/langflow/components/embeddings/util/AIMLEmbeddingsImpl.py

@@ -1,12 +1,13 @@
+import concurrent.futures
 import json
 from typing import List
 
 import httpx
-from langflow.field_typing import Embeddings
-from langchain_core.runnables.config import run_in_executor
 from langchain_core.pydantic_v1 import BaseModel, SecretStr
 from loguru import logger
 
+from langflow.field_typing import Embeddings
+
 
 class AIMLEmbeddingsImpl(BaseModel, Embeddings):
     embeddings_completion_url: str = "https://api.aimlapi.com/v1/embeddings"
@@ -15,56 +16,47 @@ class AIMLEmbeddingsImpl(BaseModel, Embeddings):
     model: str
 
     def embed_documents(self, texts: List[str]) -> List[List[float]]:
-        result_vectors = []
+        embeddings = [None] * len(texts)
-        for text in texts:
-            vector = self.embed_query(text)
-            result_vectors.append(vector)
-
-        return result_vectors
-
-    def embed_query(self, text: str) -> List[float]:
         headers = {
             "Content-Type": "application/json",
             "Authorization": f"Bearer {self.api_key.get_secret_value()}",
         }
 
+        with httpx.Client() as client:
+            with concurrent.futures.ThreadPoolExecutor() as executor:
+                futures = []
+                for i, text in enumerate(texts):
+                    futures.append((i, executor.submit(self._embed_text, client, headers, text)))
+
+                for index, future in futures:
+                    try:
+                        result_data = future.result()
+                        assert len(result_data["data"]) == 1, "Expected one embedding"
+                        embeddings[index] = result_data["data"][0]["embedding"]
+                    except (
+                        httpx.HTTPStatusError,
+                        httpx.RequestError,
+                        json.JSONDecodeError,
+                        KeyError,
+                    ) as e:
+                        logger.error(f"Error occurred: {e}")
+                        raise
+
+        return embeddings  # type: ignore
+
+    def _embed_text(self, client: httpx.Client, headers: dict, text: str) -> dict:
         payload = {
             "model": self.model,
             "input": text,
         }
-        vector = []
+        response = client.post(
-        try:
+            self.embeddings_completion_url,
-            response = httpx.post(
+            headers=headers,
-                self.embeddings_completion_url,
+            json=payload,
-                headers=headers,
+        )
-                json=payload,
+        response.raise_for_status()
-            )
+        result_data = response.json()
-            try:
+        return result_data
-                response.raise_for_status()
-                result_data = response.json()
-                vector = result_data["data"][0]["embedding"]
-            except httpx.HTTPStatusError as http_err:
-                logger.error(f"HTTP error occurred: {http_err}")
-                raise http_err
-            except httpx.RequestError as req_err:
-                logger.error(f"Request error occurred: {req_err}")
-                raise req_err
-            except json.JSONDecodeError:
-                logger.warning(f"Failed to decode JSON, response text: {response.text}")
-            except KeyError as key_err:
-                logger.warning(f"Key error: {key_err}, response content: {result_data}")
-                raise key_err
-        except httpx.TimeoutException:
-            logger.error("Request timed out.")
-            raise
-        except Exception as exc:
-            logger.error(f"Error: {exc}")
-            raise
 
-        return vector
+    def embed_query(self, text: str) -> List[float]:
-
+        return self.embed_documents([text])[0]
-    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
-        return await run_in_executor(None, self.embed_documents, texts)
-
-    async def aembed_query(self, text: str) -> List[float]:
-        return await run_in_executor(None, self.embed_query, text)

src/backend/base/langflow/components/models/AIMLModel.py

@@ -1,6 +1,6 @@
 import json
 import httpx
-from langflow.base.models.aiml_constants import CHAT_MODELS
+from langflow.base.models.aiml_constants import AIML_CHAT_MODELS
 from langflow.custom.custom_component.component import Component
 
 from langflow.inputs.inputs import FloatInput, IntInput, MessageInput, SecretStrInput
@@ -29,7 +29,7 @@ class AIMLModelComponent(Component):
         DropdownInput(
             name="model_name",
             display_name="Model Name",
-            options=CHAT_MODELS,
+            options=AIML_CHAT_MODELS,
             required=True,
         ),
         SecretStrInput(

src/backend/base/langflow/components/models/OpenAIModel.py

@@ -5,7 +5,7 @@ from langchain_openai import ChatOpenAI
 from pydantic.v1 import SecretStr
 
 from langflow.base.models.model import LCModelComponent
-from langflow.base.models.openai_constants import MODEL_NAMES
+from langflow.base.models.openai_constants import OPENAI_MODEL_NAMES
 from langflow.field_typing import LanguageModel
 from langflow.inputs import (
     BoolInput,
@@ -46,7 +46,11 @@ class OpenAIModelComponent(LCModelComponent):
             info="The schema for the Output of the model. You must pass the word JSON in the prompt. If left blank, JSON mode will be disabled.",
         ),
         DropdownInput(
-            name="model_name", display_name="Model Name", advanced=False, options=MODEL_NAMES, value=MODEL_NAMES[0]
+            name="model_name",
+            display_name="Model Name",
+            advanced=False,
+            options=OPENAI_MODEL_NAMES,
+            value=OPENAI_MODEL_NAMES[0],
         ),
         StrInput(
             name="openai_api_base",