diff --git a/src/backend/langflow/components/models/AmazonBedrock.py b/src/backend/langflow/components/models/AmazonBedrock.py new file mode 100644 index 000000000..7b38cfa52 --- /dev/null +++ b/src/backend/langflow/components/models/AmazonBedrock.py @@ -0,0 +1,62 @@ +from typing import Optional + +from langchain.llms.base import BaseLLM +from langchain.llms.bedrock import Bedrock +from langflow import CustomComponent + + +class AmazonBedrockComponent(CustomComponent): + display_name: str = "Amazon Bedrock model" + description: str = "LLM model from Amazon Bedrock." + + def build_config(self): + return { + "model_id": { + "display_name": "Model Id", + "options": [ + "ai21.j2-grande-instruct", + "ai21.j2-jumbo-instruct", + "ai21.j2-mid", + "ai21.j2-mid-v1", + "ai21.j2-ultra", + "ai21.j2-ultra-v1", + "anthropic.claude-instant-v1", + "anthropic.claude-v1", + "anthropic.claude-v2", + "cohere.command-text-v14", + ], + }, + "credentials_profile_name": {"display_name": "Credentials Profile Name"}, + "streaming": {"display_name": "Streaming", "field_type": "bool"}, + "endpoint_url": {"display_name": "Endpoint URL"}, + "region_name": {"display_name": "Region Name"}, + "model_kwargs": {"display_name": "Model Kwargs"}, + "cache": {"display_name": "Cache"}, + "code": {"advanced": True}, + "inputs": {"display_name": "Input"}, + } + + def build( + self, + inputs: str, + model_id: str = "anthropic.claude-instant-v1", + credentials_profile_name: Optional[str] = None, + region_name: Optional[str] = None, + model_kwargs: Optional[dict] = None, + endpoint_url: Optional[str] = None, + streaming: bool = False, + cache: Optional[bool] = None, + ) -> BaseLLM: + try: + output = Bedrock( + credentials_profile_name=credentials_profile_name, + model_id=model_id, + region_name=region_name, + model_kwargs=model_kwargs, + endpoint_url=endpoint_url, + streaming=streaming, + cache=cache, + ) # type: ignore + except Exception as e: + raise ValueError("Could not connect to AmazonBedrock API.") from e + return output.invoke(input=inputs) \ No newline at end of file