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)
pull/8923/head
Marie-Philippe Gill 1 year ago committed by GitHub
parent 6116cbf0de
commit 6b9f266837
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -304,6 +304,9 @@ class AmazonKendraRetriever(BaseRetriever):
client: boto3 client for Kendra client: boto3 client for Kendra
user_context: Provides information about the user context
See: https://docs.aws.amazon.com/kendra/latest/APIReference
Example: Example:
.. code-block:: python .. code-block:: python
@ -320,6 +323,7 @@ class AmazonKendraRetriever(BaseRetriever):
attribute_filter: Optional[Dict] = None attribute_filter: Optional[Dict] = None
page_content_formatter: Callable[[ResultItem], str] = combined_text page_content_formatter: Callable[[ResultItem], str] = combined_text
client: Any client: Any
user_context: Optional[Dict] = None
@validator("top_k") @validator("top_k")
def validate_top_k(cls, value: int) -> int: def validate_top_k(cls, value: int) -> int:
@ -368,6 +372,8 @@ class AmazonKendraRetriever(BaseRetriever):
} }
if self.attribute_filter is not None: if self.attribute_filter is not None:
kendra_kwargs["AttributeFilter"] = self.attribute_filter 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) response = self.client.retrieve(**kendra_kwargs)
r_result = RetrieveResult.parse_obj(response) r_result = RetrieveResult.parse_obj(response)

Loading…
Cancel
Save