refactor: AzureOpenAIEmbeddingsComponent to use new Inputs/Outputs format

This commit is contained in:
Cezar Vasconcelos 2024-06-19 21:09:42 +00:00
commit 7884dcff8a

View file

@ -1,16 +1,16 @@
from typing import Optional from typing import Optional
from langchain_core.embeddings import Embeddings
from langchain_openai import AzureOpenAIEmbeddings from langchain_openai import AzureOpenAIEmbeddings
from pydantic.v1 import SecretStr from pydantic.v1 import SecretStr
from langflow.custom import CustomComponent from langflow.base.models.model import LCModelComponent
from langflow.field_typing import Embeddings
from langflow.io import DropdownInput, IntInput, Output, SecretStrInput, TextInput
class AzureOpenAIEmbeddingsComponent(CustomComponent): class AzureOpenAIEmbeddingsComponent(LCModelComponent):
display_name: str = "Azure OpenAI Embeddings" display_name: str = "Azure OpenAI Embeddings"
description: str = "Generate embeddings using Azure OpenAI models." description: str = "Generate embeddings using Azure OpenAI models."
documentation: str = "https://python.langchain.com/docs/integrations/text_embedding/azureopenai" documentation: str = "https://python.langchain.com/docs/integrations/text_embedding/azureopenai"
beta = False
icon = "Azure" icon = "Azure"
API_VERSION_OPTIONS = [ API_VERSION_OPTIONS = [
@ -22,57 +22,56 @@ class AzureOpenAIEmbeddingsComponent(CustomComponent):
"2023-08-01-preview", "2023-08-01-preview",
] ]
def build_config(self): inputs = [
return { TextInput(
"azure_endpoint": { name="azure_endpoint",
"display_name": "Azure Endpoint", display_name="Azure Endpoint",
"required": True, required=True,
"info": "Your Azure endpoint, including the resource.. Example: `https://example-resource.azure.openai.com/`", info="Your Azure endpoint, including the resource. Example: `https://example-resource.azure.openai.com/`",
}, ),
"azure_deployment": { TextInput(
"display_name": "Deployment Name", name="azure_deployment",
"required": True, display_name="Deployment Name",
}, required=True,
"api_version": { ),
"display_name": "API Version", DropdownInput(
"options": self.API_VERSION_OPTIONS, name="api_version",
"value": self.API_VERSION_OPTIONS[-1], display_name="API Version",
"advanced": True, options=API_VERSION_OPTIONS,
}, value=API_VERSION_OPTIONS[-1],
"api_key": { advanced=True,
"display_name": "API Key", ),
"required": True, SecretStrInput(
"password": True, name="api_key",
}, display_name="API Key",
"code": {"show": False}, required=True,
"dimensions": { ),
"display_name": "Dimensions", IntInput(
"info": "The number of dimensions the resulting output embeddings should have. Only supported by certain models.", name="dimensions",
"advanced": True, display_name="Dimensions",
}, info="The number of dimensions the resulting output embeddings should have. Only supported by certain models.",
} advanced=True,
),
]
outputs = [
Output(display_name="Embeddings", name="embeddings", method="build_embeddings"),
]
def build_embeddings(self) -> Embeddings:
if not self.api_key:
raise ValueError("API Key is required")
azure_api_key = SecretStr(self.api_key)
def build(
self,
azure_endpoint: str,
azure_deployment: str,
api_version: str,
api_key: str,
dimensions: Optional[int] = None,
) -> Embeddings:
if api_key:
azure_api_key = SecretStr(api_key)
else:
azure_api_key = None
try: try:
embeddings = AzureOpenAIEmbeddings( embeddings = AzureOpenAIEmbeddings(
azure_endpoint=azure_endpoint, azure_endpoint=self.azure_endpoint,
azure_deployment=azure_deployment, azure_deployment=self.azure_deployment,
api_version=api_version, api_version=self.api_version,
api_key=azure_api_key, api_key=azure_api_key,
dimensions=dimensions, dimensions=self.dimensions,
) )
except Exception as e: except Exception as e:
raise ValueError("Could not connect to AzureOpenAIEmbeddings API.") from e raise ValueError("Could not connect to AzureOpenAIEmbeddings API.") from e