From 6b9f26683784ab6b3a7b6f41a5b40619a8a81ad7 Mon Sep 17 00:00:00 2001 From: Marie-Philippe Gill Date: Tue, 8 Aug 2023 11:37:03 -0400 Subject: [PATCH] Add user_context to AmazonKendraRetriever (#8869) ### Description Now, we can pass information like a JWT token using user_context: ```python self.retriever = AmazonKendraRetriever(index_id=kendraIndexId, user_context={"Token": jwt_token}) ``` - [x] `make lint` - [x] `make format` - [x] `make test` Also tested by pip installing in my own project, and it allows access through the token. ### Maintainers @rlancemartin, @eyurtsev ### My twitter handle [girlknowstech](https://twitter.com/girlknowstech) --- libs/langchain/langchain/retrievers/kendra.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/libs/langchain/langchain/retrievers/kendra.py b/libs/langchain/langchain/retrievers/kendra.py index 344a2bf121..c1656a3a6c 100644 --- a/libs/langchain/langchain/retrievers/kendra.py +++ b/libs/langchain/langchain/retrievers/kendra.py @@ -304,6 +304,9 @@ class AmazonKendraRetriever(BaseRetriever): client: boto3 client for Kendra + user_context: Provides information about the user context + See: https://docs.aws.amazon.com/kendra/latest/APIReference + Example: .. code-block:: python @@ -320,6 +323,7 @@ class AmazonKendraRetriever(BaseRetriever): attribute_filter: Optional[Dict] = None page_content_formatter: Callable[[ResultItem], str] = combined_text client: Any + user_context: Optional[Dict] = None @validator("top_k") def validate_top_k(cls, value: int) -> int: @@ -368,6 +372,8 @@ class AmazonKendraRetriever(BaseRetriever): } if self.attribute_filter is not None: kendra_kwargs["AttributeFilter"] = self.attribute_filter + if self.user_context is not None: + kendra_kwargs["UserContext"] = self.user_context response = self.client.retrieve(**kendra_kwargs) r_result = RetrieveResult.parse_obj(response)