[feat] add embedding dimensions parameters to openai models (#2146)

Add embedding dimensions parameters to openai models
This commit is contained in:
Jordan Frazier 2024-06-12 21:05:05 -07:00 committed by GitHub
commit 061809a94e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 15 additions and 0 deletions

View file

@ -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:

View file

@ -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,
)