Spaces:
Running
Running
| from typing import cast | |
| from langchain_community.retrievers import AmazonKendraRetriever | |
| from langflow.custom import CustomComponent | |
| from langflow.field_typing import Retriever | |
| class AmazonKendraRetrieverComponent(CustomComponent): | |
| display_name: str = "Amazon Kendra Retriever" | |
| description: str = "Retriever that uses the Amazon Kendra API." | |
| name = "AmazonKendra" | |
| icon = "Amazon" | |
| legacy: bool = True | |
| 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": { | |
| "display_name": "Attribute Filter", | |
| "field_type": "code", | |
| }, | |
| "top_k": {"display_name": "Top K", "field_type": "int"}, | |
| "user_context": { | |
| "display_name": "User Context", | |
| "field_type": "code", | |
| }, | |
| "code": {"show": False}, | |
| } | |
| def build( | |
| self, | |
| index_id: str, | |
| top_k: int = 3, | |
| region_name: str | None = None, | |
| credentials_profile_name: str | None = None, | |
| attribute_filter: dict | None = None, | |
| user_context: dict | None = None, | |
| ) -> Retriever: # type: ignore[type-var] | |
| try: | |
| 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, | |
| ) | |
| except Exception as e: | |
| msg = "Could not connect to AmazonKendra API." | |
| raise ValueError(msg) from e | |
| return cast("Retriever", output) | |