Amazon Bedrock Knowledge Base – Part 3 (SDK)

Since the last writeup, AWS has added support for Anthropic Claude 3 model to AWS Bedrock Knowledge Base (ABKB). It has also added the ability to add your own metadata to your source files in order to perform filtering when doing query. For example, you may want to add a metadata to a certain set of files to indicate that they are from year 2023. Then during your query, you can include a filter to indicate you only want to use data from year 2023. This provides another set of tools for developers to create more relevant and targeted query. Note that filtering is only supported for FAISS vector engine.

If you’re looking to integrate ABKB into your code, there are two primary methods: using one of the AWS SDK or interacting through HTTP API. In this article, we will be using Boto3, the AWS SDK for Python. Here is a simple example to do a retrieve and generate query using Boto3. This example uses the new Claude 3 Sonnet model.

import boto3
import json

AWS_ACCESS_KEY="_your_access_key_"
AWS_SECRET_KEY="_your_secret_key_"
REGION_NAME="_your_region_"

client = boto3.client('bedrock-agent-runtime',
                      aws_access_key_id=AWS_ACCESS_KEY,
                      aws_secret_access_key=AWS_SECRET_KEY,
                      region_name=REGION_NAME
)

# retrieval and generate
response = client.retrieve_and_generate(
    input={
        'text': 'how to apply for leave'
    },
    retrieveAndGenerateConfiguration={
        'knowledgeBaseConfiguration': {
            'knowledgeBaseId': 'LEBQPJQ9BY',
            'modelArn': 'arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-3-sonnet-20240229-v1:0',
            'retrievalConfiguration': {
                'vectorSearchConfiguration': {
                    'overrideSearchType': 'HYBRID'
                }
            }
        },
        'type': 'KNOWLEDGE_BASE'
    }
)

print(json.dumps(response))

Running the code produces the following output in JSON:

{
  "ResponseMetadata": {
		...trimmed...
  },
  "citations": [
    {
      "generatedResponsePart": {
        "textResponsePart": {
          "span": {
            "end": 705,
            "start": 364
          },
          "text": "...trimmed..."
        }
      },
      "retrievedReferences": [
        {
          "content": {
            "text": "...trimmed..."
          },
          "location": {
            "s3Location": {
              "uri": "s3://...trimmed..."
            },
            "type": "S3"
          }
        }
      ]
    }
  ],
  "output": {
    "text": "To apply for leave as an employee on the Workday mobile app:\n\n1. Navigate to your Workday Mobile Homepage and select 'View Applications' under 'Frequently Used'\n2. Select 'Time Off'\n3. Select the date(s) you want to apply for leave\n4. Select 'Next' and choose the leave type\n5. Select any required reasons or upload attachments if applicable\n6. Submit the request To apply for leave as an employee on the Workday desktop:\n\n1. Go to the Workday Homepage and select the 'Absence' worklet\n2. Under 'Request', select 'Request Absence'\n3. Select the date(s) for the leave and click 'Request Absence'\n4. Choose the leave type\n5. Select 'Next' and provide any required reasons or attachments\n6. Submit the request"
  },
  "sessionId": "c8332417-df3c-41e5-8516-ad38cc09de15"
}

For this simple task there is not much difference in output from the various Claude models. I expect the differences will be more pronounced for complex tasks or those involving much larger context window.

With this, we conclude the three-part series on Amazon Bedrock Knowledge Base. I have covered everything from creating the knowledge base, testing it in the playground, to executing queries via CLI and SDK. Hopefully this gives a good overview of the processes involved and capabilities of this new service.