diff --git a/src/backend/base/langflow/base/models/aiml_constants.py b/src/backend/base/langflow/base/models/aiml_constants.py index 41955a949..a96219010 100644 --- a/src/backend/base/langflow/base/models/aiml_constants.py +++ b/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", diff --git a/src/backend/base/langflow/base/models/openai_constants.py b/src/backend/base/langflow/base/models/openai_constants.py index 474294ced..7488aa52a 100644 --- a/src/backend/base/langflow/base/models/openai_constants.py +++ b/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", +] diff --git a/src/backend/base/langflow/components/embeddings/AIMLEmbeddings.py b/src/backend/base/langflow/components/embeddings/AIMLEmbeddings.py index 3bd1d3cf9..e52198ba4 100644 --- a/src/backend/base/langflow/components/embeddings/AIMLEmbeddings.py +++ b/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( diff --git a/src/backend/base/langflow/components/embeddings/OpenAIEmbeddings.py b/src/backend/base/langflow/components/embeddings/OpenAIEmbeddings.py index ef810a65a..f37ee7641 100644 --- a/src/backend/base/langflow/components/embeddings/OpenAIEmbeddings.py +++ b/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=[ - "text-embedding-3-small", - "text-embedding-3-large", - "text-embedding-ada-002", - ], + options=OPENAI_EMBEDDING_MODEL_NAMES, value="text-embedding-3-small", ), DictInput(name="model_kwargs", display_name="Model Kwargs", advanced=True), diff --git a/src/backend/base/langflow/components/embeddings/util/AIMLEmbeddingsImpl.py b/src/backend/base/langflow/components/embeddings/util/AIMLEmbeddingsImpl.py index 74bf2110c..34859dd81 100644 --- a/src/backend/base/langflow/components/embeddings/util/AIMLEmbeddingsImpl.py +++ b/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 = [] - for text in texts: - vector = self.embed_query(text) - result_vectors.append(vector) - - return result_vectors - - def embed_query(self, text: str) -> List[float]: + embeddings = [None] * len(texts) 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 = [] - try: - response = httpx.post( - self.embeddings_completion_url, - headers=headers, - json=payload, - ) - try: - 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 + response = client.post( + self.embeddings_completion_url, + headers=headers, + json=payload, + ) + response.raise_for_status() + result_data = response.json() + return result_data - return vector - - 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) + def embed_query(self, text: str) -> List[float]: + return self.embed_documents([text])[0] diff --git a/src/backend/base/langflow/components/models/AIMLModel.py b/src/backend/base/langflow/components/models/AIMLModel.py index 37233860f..df0904590 100644 --- a/src/backend/base/langflow/components/models/AIMLModel.py +++ b/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( diff --git a/src/backend/base/langflow/components/models/OpenAIModel.py b/src/backend/base/langflow/components/models/OpenAIModel.py index 355ab567b..d621bbb9a 100644 --- a/src/backend/base/langflow/components/models/OpenAIModel.py +++ b/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",