From 31245c2d964ac3fff7c98eabe197397a39592f06 Mon Sep 17 00:00:00 2001
From: kandakji <kondakji.oussama@gmail.com>
Date: Fri, 27 Oct 2023 08:14:48 +0200
Subject: [PATCH] Bedrock Embeddings custom component

---
 .../embeddings/AmazonBedrockEmbeddings.py     | 46 +++++++++++++++++++
 .../components/embeddings/__init__.py         |  0
 src/backend/langflow/config.yaml              |  2 +-
 src/frontend/src/utils/styleUtils.ts          |  2 +-
 4 files changed, 48 insertions(+), 2 deletions(-)
 create mode 100644 src/backend/langflow/components/embeddings/AmazonBedrockEmbeddings.py
 create mode 100644 src/backend/langflow/components/embeddings/__init__.py

diff --git a/src/backend/langflow/components/embeddings/AmazonBedrockEmbeddings.py b/src/backend/langflow/components/embeddings/AmazonBedrockEmbeddings.py
new file mode 100644
index 000000000..6efa22961
--- /dev/null
+++ b/src/backend/langflow/components/embeddings/AmazonBedrockEmbeddings.py
@@ -0,0 +1,46 @@
+from typing import Optional
+from langflow import CustomComponent
+
+from langchain.embeddings import BedrockEmbeddings
+from langchain.embeddings.base import Embeddings
+
+
+class AmazonBedrockEmeddingsComponent(CustomComponent):
+    """
+    A custom component for implementing an Embeddings Model using Amazon Bedrock.
+    """
+
+    display_name: str = "Amazon Bedrock Embeddings"
+    description: str = "Embeddings model from Amazon Bedrock."
+    documentation = "https://python.langchain.com/docs/modules/data_connection/text_embedding/integrations/bedrock"
+    beta = True
+
+    def build_config(self):
+        return {
+            "model_id": {
+                "display_name": "Model Id",
+                "options": ["amazon.titan-embed-text-v1"],
+            },
+            "credentials_profile_name": {"display_name": "Credentials Profile Name"},
+            "endpoint_url": {"display_name": "Bedrock Endpoint URL"},
+            "region_name": {"display_name": "AWS Region"},
+            "code": {"show": False},
+        }
+
+    def build(
+        self,
+        model_id: str = "amazon.titan-embed-text-v1",
+        credentials_profile_name: Optional[str] = None,
+        endpoint_url: Optional[str] = None,
+        region_name: Optional[str] = None,
+    ) -> Embeddings:
+        try:
+            output = BedrockEmbeddings(
+                credentials_profile_name=credentials_profile_name,
+                model_id=model_id,
+                endpoint_url=endpoint_url,
+                region_name=region_name,
+            )  # type: ignore
+        except Exception as e:
+            raise ValueError("Could not connect to AmazonBedrock API.") from e
+        return output
diff --git a/src/backend/langflow/components/embeddings/__init__.py b/src/backend/langflow/components/embeddings/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/backend/langflow/config.yaml b/src/backend/langflow/config.yaml
index cfa09db77..efcdc69cc 100644
--- a/src/backend/langflow/config.yaml
+++ b/src/backend/langflow/config.yaml
@@ -106,7 +106,7 @@ embeddings:
     documentation: "https://python.langchain.com/docs/modules/data_connection/text_embedding/integrations/cohere"
   VertexAIEmbeddings:
     documentation: "https://python.langchain.com/docs/modules/data_connection/text_embedding/integrations/google_vertex_ai_palm"
-  BedrockEmbeddings:
+  AmazonBedrockEmbeddings:
     documentation: "https://python.langchain.com/docs/modules/data_connection/text_embedding/integrations/bedrock"
 
 llms:
diff --git a/src/frontend/src/utils/styleUtils.ts b/src/frontend/src/utils/styleUtils.ts
index abcb27d1c..86011f6e0 100644
--- a/src/frontend/src/utils/styleUtils.ts
+++ b/src/frontend/src/utils/styleUtils.ts
@@ -200,7 +200,7 @@ export const nodeIconsLucide: iconsType = {
   ArrowUpToLine: ArrowUpToLine,
   Chroma: ChromaIcon,
   AirbyteJSONLoader: AirbyteIcon,
-  BedrockEmbeddings: AWSIcon,
+  AmazonBedrockEmbeddings: AWSIcon,
   Anthropic: AnthropicIcon,
   ChatAnthropic: AnthropicIcon,
   BingSearchAPIWrapper: BingIcon,