Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adding bedrock metadata filter and search type parameters #104

Merged
merged 10 commits into from
Jul 17, 2024

Conversation

ravediamond
Copy link
Contributor

Description

This pull request enhances the functionality of theAmazonKnowledgeBasesRetriever class by adding optional parameters for the following features (already implemented in boto3):

  • added filter parameter in VectorSearchConfig to filter on metadata tags.
  • added overrideSearchType in VectorSearchConfig to override the default search type to be able to choose between 'Default', 'Hybrid' or 'Semantic'

Here an example answer before the change:

retriever = AmazonKnowledgeBasesRetriever(
    knowledge_base_id="KNOWLEDGE_BASE_ID,
    retrieval_config={
        "vectorSearchConfiguration": {
            "numberOfResults": 4
        }
    },
)

And now after the change:

retriever = AmazonKnowledgeBasesRetriever(
    knowledge_base_id="KNOWLEDGE_BASE_ID,
    retrieval_config={
        "vectorSearchConfiguration": {
            "numberOfResults": 4,
            "filter": {
                "equals": {"key": "type", "TAG1"}
            },
            "overrideSearchType": "HYBRID"
        }
    },
)

Change

  • added SearchFilter class to handle parameter filtering
  • added filter parameter in VectorSearchConfig
  • added overrideSearchType in VectorSearchConfig

Motivation

This will allow users to better use the full capacity of Bedrock Knowledge Base.

Additional Notes

  • This change maintains backwards compatibility with existing usage.
  • The nextToken feature when with very long results has not been implemented.

@ravediamond
Copy link
Contributor Author

@3coins can you please have a look? I confirmed that this PR is working.
Thanks in advance.

Copy link
Collaborator

@3coins 3coins left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ravediamond
Thanks for making these changes. Added some suggestions to simplify the code and handle the nextToken.

libs/aws/langchain_aws/retrievers/bedrock.py Outdated Show resolved Hide resolved
libs/aws/langchain_aws/retrievers/bedrock.py Outdated Show resolved Hide resolved
libs/aws/langchain_aws/retrievers/bedrock.py Outdated Show resolved Hide resolved
@ravediamond
Copy link
Contributor Author

@3coins Thank you for the review. I correct my PR as per your recommendations. What do you think ?

libs/aws/langchain_aws/retrievers/bedrock.py Outdated Show resolved Hide resolved
libs/aws/langchain_aws/retrievers/bedrock.py Outdated Show resolved Hide resolved
libs/aws/langchain_aws/retrievers/bedrock.py Outdated Show resolved Hide resolved
libs/aws/langchain_aws/retrievers/bedrock.py Outdated Show resolved Hide resolved
@ravediamond ravediamond requested a review from 3coins July 12, 2024 14:43
@ravediamond
Copy link
Contributor Author

Ok I corrected as per your reviews and I will create another PR for the nextToken when it is merged.
Thanks in advance fir the review @3coins and also for having merged my other PR.

Copy link
Collaborator

@3coins 3coins left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ravediamond
Made a small update to exclude None values from filter, this was failing the integration test.
Looks good now, thanks for making all the updates.

@3coins 3coins merged commit 19822bc into langchain-ai:main Jul 17, 2024
12 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants