Add argument for AmazonBedrock and Amazon Kendra

This commit is contained in:
Yuki Sekiya 2023-10-15 02:25:45 +09:00
commit 2cafb33b8e
2 changed files with 28 additions and 10 deletions

View file

@ -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

View file

@ -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