diff --git a/src/backend/base/langflow/components/embeddings/AzureOpenAIEmbeddings.py b/src/backend/base/langflow/components/embeddings/AzureOpenAIEmbeddings.py index dd40d64d5..4fca09762 100644 --- a/src/backend/base/langflow/components/embeddings/AzureOpenAIEmbeddings.py +++ b/src/backend/base/langflow/components/embeddings/AzureOpenAIEmbeddings.py @@ -1,3 +1,4 @@ +from typing import Optional from langchain_core.embeddings import Embeddings from langchain_openai import AzureOpenAIEmbeddings from pydantic.v1 import SecretStr @@ -44,6 +45,11 @@ class AzureOpenAIEmbeddingsComponent(CustomComponent): "password": True, }, "code": {"show": False}, + "dimensions": { + "display_name": "Dimensions", + "info": "The number of dimensions the resulting output embeddings should have. Only supported by certain models.", + "advanced": True, + }, } def build( @@ -52,6 +58,7 @@ class AzureOpenAIEmbeddingsComponent(CustomComponent): azure_deployment: str, api_version: str, api_key: str, + dimensions: Optional[int] = None, ) -> Embeddings: if api_key: azure_api_key = SecretStr(api_key) @@ -63,6 +70,7 @@ class AzureOpenAIEmbeddingsComponent(CustomComponent): azure_deployment=azure_deployment, api_version=api_version, api_key=azure_api_key, + dimensions=dimensions, ) except Exception as e: diff --git a/src/backend/base/langflow/components/embeddings/OpenAIEmbeddings.py b/src/backend/base/langflow/components/embeddings/OpenAIEmbeddings.py index 2ff78d562..4813fc793 100644 --- a/src/backend/base/langflow/components/embeddings/OpenAIEmbeddings.py +++ b/src/backend/base/langflow/components/embeddings/OpenAIEmbeddings.py @@ -84,6 +84,11 @@ class OpenAIEmbeddingsComponent(CustomComponent): "advanced": True, }, "tiktoken_enable": {"display_name": "TikToken Enable", "advanced": True}, + "dimensions": { + "display_name": "Dimensions", + "info": "The number of dimensions the resulting output embeddings should have. Only supported by certain models.", + "advanced": True, + }, } def build( @@ -109,6 +114,7 @@ class OpenAIEmbeddingsComponent(CustomComponent): skip_empty: bool = False, tiktoken_enable: bool = True, tiktoken_model_name: Optional[str] = None, + dimensions: Optional[int] = None, ) -> Embeddings: # This is to avoid errors with Vector Stores (e.g Chroma) if disallowed_special == ["all"]: @@ -140,4 +146,5 @@ class OpenAIEmbeddingsComponent(CustomComponent): show_progress_bar=show_progress_bar, skip_empty=skip_empty, tiktoken_model_name=tiktoken_model_name, + dimensions=dimensions, )