From 2cafb33b8eb78e403ad24646ac594bc5f8a4578c Mon Sep 17 00:00:00 2001
From: Yuki Sekiya <ysekiy@amazon.co.jp>
Date: Sun, 15 Oct 2023 02:25:45 +0900
Subject: [PATCH] Add argument for AmazonBedrock and Amazon Kendra

---
 .../langflow/components/llms/AmazonBedrock.py | 14 +++++------
 .../components/retrievers/AmazonKendra.py     | 24 +++++++++++++++++--
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/src/backend/langflow/components/llms/AmazonBedrock.py b/src/backend/langflow/components/llms/AmazonBedrock.py
index 38fc5ff3e..c829ef073 100644
--- a/src/backend/langflow/components/llms/AmazonBedrock.py
+++ b/src/backend/langflow/components/llms/AmazonBedrock.py
@@ -10,7 +10,6 @@ class AmazonBedrockComponent(CustomComponent):
 
     def build_config(self):
         return {
-            "credentials_profile_name": {"display_name": "Credentials Profile Name", "password": True},
             "model_id": {
                 "display_name": "Model Id",
                 "options": [
@@ -26,25 +25,24 @@ class AmazonBedrockComponent(CustomComponent):
                     "cohere.command-text-v14",
                 ],
             },
-            "model_kwargs": {
-                "display_name": "Model Keyword Arguments",
-                "field_type": "code",
+            "credentials_profile_name": {"display_name": "Credentials Profile Name"},
+            "streaming": {
+                "display_name": "Streaming",
+                "field_type": "bool"
             },
             "code": {"show": False},
         }
 
     def build(
         self,
-        credentials_profile_name: str,
         model_id: str = "anthropic.claude-instant-v1",
-        model_kwargs: Optional[dict] = None,
+        credentials_profile_name: Optional[str] = None,
     ) -> BaseLLM:
         try:
             output = Bedrock(
                 credentials_profile_name=credentials_profile_name,
                 model_id=model_id,
-                model_kwargs=model_kwargs,
-            )
+            )  # 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/retrievers/AmazonKendra.py b/src/backend/langflow/components/retrievers/AmazonKendra.py
index a964847df..72ac96757 100644
--- a/src/backend/langflow/components/retrievers/AmazonKendra.py
+++ b/src/backend/langflow/components/retrievers/AmazonKendra.py
@@ -1,3 +1,4 @@
+from typing import Optional
 from langflow import CustomComponent
 from langchain.retrievers import AmazonKendraRetriever
 from langchain.schema import BaseRetriever
@@ -10,18 +11,37 @@ class AmazonKendraRetrieverComponent(CustomComponent):
     def build_config(self):
         return {
             "index_id": {"display_name": "Index ID"},
+            "region_name": {"display_name": "Region Name"},
+            "credentials_profile_name": {"display_name": "Credentials Profile Name"},
             "attribute_filter": {
                 "attribute_filter": "Attribute Filter",
                 "field_type": "code",
             },
+            "top_k": {"display_name": "Top K", "field_type": "int"},
+            "user_context": {
+                "attribute_filter": "User Context",
+                "field_type": "code",
+            },
             "code": {"show": False},
         }
 
     def build(
-        self, index_id: str, attribute_filter: dict
+        self,
+        index_id: str,
+        top_k: int = 3,
+        region_name: Optional[str] = None,
+        credentials_profile_name: Optional[str] = None,
+        attribute_filter: Optional[dict] = None,
+        user_context: Optional[dict] = None,
     ) -> BaseRetriever:
         try:
-            output = AmazonKendraRetriever(index_id=index_id, attribute_filter=attribute_filter)
+            output = AmazonKendraRetriever(
+                index_id=index_id,
+                top_k=top_k,
+                region_name=region_name,
+                credentials_profile_name=credentials_profile_name,
+                attribute_filter=attribute_filter,
+                user_context=user_context)  # type: ignore
         except Exception as e:
             raise ValueError("Could not connect to AmazonKendra API.") from e
         return output