Knowledge Base Support for the Generic Bedrock Agent Test UI
Published: 2024-05-21 19:42:09

## Introduction

In the blog post [Developing a Generic Streamlit UI to Test Amazon Bedrock Agents](https://dev.to/aws-builders/developing-a-generic-streamlit-ui-to-test-amazon-bedrock-agents-5c7o), I shared the design and [source code](https://github.com/acwwat/amazon-bedrock-agent-test-ui) of a basic yet functional UI for testing Bedrock agents. Since then, I've explored Knowledge Bases for Amazon Bedrock and shared my insights in another blog post, [Adding an Amazon Bedrock Knowledge Base to the Forex Rate Assistant](https://dev.to/aws-builders/building-a-basic-forex-rate-assistant-using-agents-for-amazon-bedrock-17fp). If you haven't checked it out yet, I highly recommend doing so.

The Bedrock console offers additional features for testing agents that are integrated with knowledge bases, including citations in the responses and trace information about the retrieved results from knowledge bases:


![Knowledge base citations and traces](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/8g75b1me07aka9sagumm.png)

With a bit of work, I have added similar support to the generic test UI and I am happy to share the updates in the [GitHub repository](https://github.com/acwwat/amazon-bedrock-agent-test-ui).

## Design overview

With the latest update, citations are now added to the response in a manner similar to how they are displayed in the Bedrock console:

![Citations in the response](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/o9yo9v7pm7c1xe6czk3g.png)

The [Agent for Bedrock Runtime API](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_InvokeAgent.html) provides, for each citation, the start and end index of the text that references the source as well as the document location in S3. Through string manipulation, citation numbers are incorporated into the response text, and references are appended to the end of the text. Spacing is a bit finicky due to the use of [markdown](https://docs.streamlit.io/develop/api-reference/text/st.markdown), so some HTML markups are used.

Another feature is the inclusion of citation details in the **Trace** section of the left pane:

![Citation details](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/ckfs00yw4wp3liu6t67c.png)

Each citation block provides the raw output from the API response and includes the information mentioned earlier, as well as the raw results retrieved by the knowledge base.

Additionally, the trace blocks from the original test UI can now provide more detailed information for knowledge base invocations. You will find all references retrieved from the knowledge base, which form a superset of the citations in the final response after the model processes them.

## Summary

With the improvements to the generic test UI outlined in this post, you should now be able to test any Bedrock agents with attached knowledge bases. I hope you find this update helpful and look forward to further enhancements as I continue on my Amazon Bedrock journey.

Be sure to explore other posts on the [Avangards Blog](https://blog.avangards.io) to learn more about generative AI in AWS, Terraform, and other technical topics. Have a great day!