jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

refactor: Update AmazonBedrockEmbeddingsComponent to use new Inputs/Outputs format

Browse source
This commit is contained in:
Cezar Vasconcelos 2024-06-19 21:09:54 +00:00
commit a180589f79
1 changed files with 35 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

62
src/backend/base/langflow/components/embeddings/AmazonBedrockEmbeddings.py
View file

@ -1,42 +1,50 @@
from typing import Optional from typing import Optional
from langchain_community.embeddings import BedrockEmbeddings from langchain_community.embeddings import BedrockEmbeddings
from langchain_core.embeddings import Embeddings from langflow.base.models.model import LCModelComponent
from langflow.field_typing import Embeddings
from langflow.custom import CustomComponent from langflow.io import DropdownInput, Output, SecretStrInput, TextInput
class AmazonBedrockEmeddingsComponent(CustomComponent): class AmazonBedrockEmbeddingsComponent(LCModelComponent):
display_name: str = "Amazon Bedrock Embeddings" display_name: str = "Amazon Bedrock Embeddings"
description: str = "Generate embeddings using Amazon Bedrock models." description: str = "Generate embeddings using Amazon Bedrock models."
documentation = "https://python.langchain.com/docs/modules/data_connection/text_embedding/integrations/bedrock" documentation = "https://python.langchain.com/docs/modules/data_connection/text_embedding/integrations/bedrock"
icon = "Amazon"
def build_config(self): inputs = [
return { DropdownInput(
"model_id": { name="model_id",
"display_name": "Model Id", display_name="Model Id",
"options": ["amazon.titan-embed-text-v1"], options=["amazon.titan-embed-text-v1"],
}, value="amazon.titan-embed-text-v1",
"credentials_profile_name": {"display_name": "Credentials Profile Name"}, ),
"endpoint_url": {"display_name": "Bedrock Endpoint URL"}, TextInput(
"region_name": {"display_name": "AWS Region"}, name="credentials_profile_name",
"code": {"show": False}, display_name="Credentials Profile Name",
} ),
TextInput(
name="endpoint_url",
display_name="Bedrock Endpoint URL",
),
TextInput(
name="region_name",
display_name="AWS Region",
),
]
def build( outputs = [
self, Output(display_name="Embeddings", name="embeddings", method="build_embeddings"),
model_id: str = "amazon.titan-embed-text-v1", ]
credentials_profile_name: Optional[str] = None,
endpoint_url: Optional[str] = None, def build_embeddings(self) -> Embeddings:
region_name: Optional[str] = None,
) -> Embeddings:
try: try:
output = BedrockEmbeddings( output = BedrockEmbeddings(
credentials_profile_name=credentials_profile_name, credentials_profile_name=self.credentials_profile_name,
model_id=model_id, model_id=self.model_id,
endpoint_url=endpoint_url, endpoint_url=self.endpoint_url,
region_name=region_name, region_name=self.region_name,
) # type: ignore ) # type: ignore
except Exception as e: except Exception as e:
raise ValueError("Could not connect to AmazonBedrock API.") from e raise ValueError("Could not connect to Amazon Bedrock API.") from e
return output return output