from typing import Optional
from langflow import CustomComponent
from langchain.retrievers import AmazonKendraRetriever
from langchain.schema import BaseRetriever


class AmazonKendraRetrieverComponent(CustomComponent):
    display_name: str = "Amazon Kendra Retriever"
    description: str = "Retriever that uses the Amazon Kendra API."

    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: 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,
                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
