Knowledge Bases in AWS Bedrock

An AWS Knowledge Base is to Bedrock as the Retrieval component is to LangChain🦜🔗. Both offer tools to perform RAG, but Knowledge Bases pair nicely with your Bedrock models, and offer a managed RAG solution that you can quickly standup in front of your foundation models. If you're not already familiar with what I'm talking about, I'll catch you up. Large language models are rad. When used correctly, the information they string together can be profoundly useful, and when given enough "context" alongside clear instruction, we've found these models are capable of integrating that "context" to a similar effect as if it was a traditional ML model trained on the same data.

the problem...

Models all have a limited context windows. These windows define the size limits of both input and output measured in tokens as a single value. Essentially this is to say it becomes impossible to include the entire contents of data lake if you need this as context for your prompts.

large context windows

It is worth calling out that as of late we've seen models boasting 100K+ token windows like Anthropic's Claude, and even Google boasting a version of Gemini with a million+ limit. In Bedrock at the time of writing this, Claude V3 Sonnet and Haiku are the largest options with 200k limits.

These large windows enable much more complex prompts without the need for augmenting. Documents that previously may have needed to be chunked, summarized, or otherwise linguistically compressed, may not need to suffer from these lossy tactics with the large enough windows (obviously at a cost in tokens...).

The Benefits of Incorporating Knowledge Bases

Retrieval Augmented Generation or RAG, is a technique that allows you to create more contextually rich responses from Large Language Models (LLMs). By drawing information from a data store, models can provide more accurate and informative responses to queries/prompts. The AWS Bedrock Knowledge Base approach speeds up development time by providing a ready-to-use solution with consoles, APIs, and tools to manage all the parts. Things like your document loaders, source and vector store, embeddings model, chunking rules, and retriever are all considered and configurable from our favorite 🐬boto3 library (if you're in Python).

Remember however that LangChain is powerful and actively maintained/growing. It does everything that a Knowledge Bases offers, but is agnostic of clouds and models. So don't write it off. If find yourself only needing "part of" or a custom solution you may find it super useful.
LangChain Retrieval: https://python.langchain.com/v0.1/docs/modules/data_connection/

Implementing Knowledge Bases in Your AWS Bedrock Projects

I'm not going to make this a whole tutorial on the setup, AWS has done a great job providing notebooks and documentation on this. Not to mention the Bedrock console provide a great introduction to what these tools do.

This graphic from the AWS Bedrocks Workshops page roughly illustrates the Knowledge Base flow (and definitely note how closely it resembles what we see from LangChain.)

The basic idea that you'll get into when working through the AWS Knowledge Base material is as follows:

1. You setup an S3 Bucket that will serve as your source storage.

2. From this, the contents will be "chunked" or essentially broken down into smaller parts, as intelligently as possible based on a strategy you set.

3. You pick or use the suggested embeddings model, which is just another LLM trained or tuned on the specific task on embeddings.

4. And lastly the vector store is by default OpenSearch, and is very easily managed via your Knowledge Base (console or API).

Okay, but what I really wanted to include was 2 tips for anyone starting out, as the information above is pretty easy to come by.

Tip 1

Consider your embeddings retention policy.

While a simple configuration, its worth thinking about, and might make you really consider how you want to use your Knowledge Base. A Retain policy will persist your vector data regardless of if the source in still available in the S3 bucket. This is great if you want to maintain an ever growing vector store without a massive S3 bucket.

But consider the Delete policy. This will remove the related vector data if the source is deleted from S3. This makes it very easy to control your embeddings, as their existence will just be based on whats in an S3 bucket.

Tip 2

And that takes me to my second tip. When you start using a Knowledge Bases you may quickly find the need to filter or limit what your retriever is allowed to return from. Meaning you may have a large pool of documents, but perhaps you know for a given prompt that you're only concerned with one type of document.

This is where metaDataAttributes and filtering upon retrieval comes in handy.
To include metadata for a file in your data source, create a JSON file consisting of a metadataAttributes field that maps to an object with a key-value pair for each metadata attribute. Then simply upload it to the same folder and bucket as the source document file.

{
   "metadataAttributes": {
      "${attribute1}": "${value1}",
      "${attribute2}": "${value2}",
      ...
   }
}

Append .metadata.json after the file extension (for example, if you have a file named Something.txt, the metadata file must be named Something.txt.metadata.json

Now when performing a retrieval, we can use these attributes to filter what sources the parts are allowed to come from.  KnowledgeBaseRetrievalConfiguration

Heres are the tables of operators that can be used to filter.

AWS has more examples in the links above, but heres 2 to give you an idea and get you started.

One filtering operator.

    "retrievalConfiguration": {
    "vectorSearchConfiguration": {
        "filter": {
            "<filter-type>": {
                "key": "string",
                "value": "string" | number | boolean | ["string", "string", ...]
            }
        }
    }
}

One logical operator.

"retrievalConfiguration": {
    "vectorSearchConfiguration": {
        "filter": {
            "andAll | orAll": [
                "<filter-type>": {
                    "key": "string",
                    "value": "string" | number | boolean | ["string", "string", ...]
                },
                "<filter-type>": {
                    "key": "string",
                    "value": "string" | number | boolean | ["string", "string", ...]
                },
                ...
            ]
        }
    }
}

Good luck, have fun.