Developing a Generic Streamlit UI to Test Amazon Bedrock Agents
Tags: aws, ai, python
Published: 2024-05-05 23:44:54

>💡 **Updated 2024-05-21:** The UI has been updated with features to support knowledge bases that are attached to the agent. For details, refer to the blog post [Knowledge Base Support for the Generic Bedrock Agent Test UI](https://dev.to/aws-builders/knowledge-base-support-for-the-generic-bedrock-agent-test-ui-4gfg).

## Introduction

In the earlier blog post [Building a Basic Forex Rate Assistant Using Agents for Amazon Bedrock](https://dev.to/aws-builders/building-a-basic-forex-rate-assistant-using-agents-for-amazon-bedrock-17fp), I walked readers through the process of building and testing a Bedrock agent in the AWS Management Console. While the built-in test interface is great for validation as changes are made in the Bedrock console, it is not scalable to other team members such as testers who often don't have direct access to AWS.

Meanwhile, a developer workflow that does not require access to AWS Management Console may provide a better experience. As a developer, I appreciate having an integrated development environment (IDE) such as [Visual Studio Code](https://code.visualstudio.com/) where I can code, deploy, and test in one place.

To address these two challenges, I decided to build a basic but functional UI for testing Bedrock agents. In this blog post, I share with readers the end product and some details about its design.

## Design and implementation overview

The following is the list of requirements that I defined for the test UI:

* The design should be minimal but functional, since the focus is not on the UI but on being able to validate the business logic of the agents.
    
* The solution must provide the basic features as the Bedrock console, including trace.
    
* The solution must be adaptable to any Bedrock agents with no to minimal changes.
    
* The solution must run both locally and as a shared webapp for different workflows.
    

I decided to use [Streamlit](https://streamlit.io/) to build the UI as it is a popular and fitting choice. Streamlit is an open-source Python library used for building interactive web applications specially for AI and data applications. Since the application code is written only in Python, it is easy to learn and build with.

The [Agents for Amazon Bedrock Runtime API](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Agents_for_Amazon_Bedrock_Runtime.html) can be used to interact with a Bedrock agent. Since the Streamlit app is developed in Python, we will naturally use the [AWS SDK for Python (Boto3)](https://boto3.amazonaws.com/v1/documentation/api/latest/index.html) for the integration. The [AWS SDK Code Examples](https://aws.amazon.com/developer/code-examples/) code library provides an [example](https://docs.aws.amazon.com/code-library/latest/ug/bedrock-agent-runtime_example_bedrock-agent-runtime_InvokeAgent_section.html) on how to use the [`AgentsforBedrockRuntime.Client.invoke_agent` function](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent-runtime/client/invoke_agent.html) to call the Bedrock agent. The function documentation was essential to determine the response format and the information.

The UI design is rather minimal as you can see in the following screenshot:

![Test UI design elements](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/ibvvnbvg7p8fgycbo95d.png)

In the left [sidebar](https://docs.streamlit.io/develop/api-reference/layout/st.sidebar), I include elements related to troubleshooting such as a session reset button and trace information similar to the Bedrock console. The main pane is a simple chat interface with the [messages](https://docs.streamlit.io/develop/api-reference/chat/st.chat_message) and the [input](https://docs.streamlit.io/develop/api-reference/chat/st.chat_input). I made the favicon and page title configurable using environment variables for a white label experience.

## About the repository structure

You can find the source code for the test UI in the [acwwat/amazon-bedrock-agent-test-ui](https://github.com/acwwat/amazon-bedrock-agent-test-ui) GitHub repository. The repository structure follows a [standard structure](https://github.com/markdouthwaite/streamlit-project) as recommended by Mark Douthwaite for Streamlit projects. For a detailed explanation of the structure, refer to the [getting started](https://github.com/markdouthwaite/streamlit-project/blob/master/docs/template-info.md) documentation. The only tweak I made is that I put the backend integration code into `bedrock_agent_runtime.py` in a `services` directory.

```plaintext
├── services
│   ├── bedrock_agent_runtime.py
├── .gitignore
├── app.py
├── Dockerfile
├── LICENSE
├── README.md
├── requirements.txt
```

## Configuring and running the app locally

To run the Streamlit app locally, you just need to have the [AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html) and [Python 3](https://www.python.org/downloads/) installed. Then you can clone the [acwwat/amazon-bedrock-agent-test-ui](https://github.com/acwwat/amazon-bedrock-agent-test-ui) GitHub repository and follow the steps below:

1. Run the following command to install the dependencies:  
    `pip install -r requirements.txt`
    
2. Configure the environment variables for the AWS CLI and Boto3. You would typically [configure the AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/cli-chap-configure.html) to create a named profile, then set the [`AWS_PROFILE` environment variable](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html#cli-configure-files-using-profiles) to refer to it.
    
3. Set the following environment variables as appropriate:
    - `BEDROCK_AGENT_ID` - The ID of the Bedrock agent, which you can find in the Bedrock console or by running the [`aws bedrock-agent list-agents`](https://docs.aws.amazon.com/cli/latest/reference/bedrock-agent/list-agents.html) command.
    - `BEDRROCK_AGENT_ALIAS_ID` - The ID of the agent alias, which you can find in the Bedrock console or by running the [`aws bedrock-agent list-agent-aliases`](https://docs.aws.amazon.com/cli/latest/reference/bedrock-agent/list-agent-aliases.html) command. If this environment variable is not set, the default test alias ID `TSTALIASID` will be used.
    - `BEDROCK_AGENT_TEST_UI_TITLE` - (Optional) The page title. If this environment is not set, the generic title in the above screenshot will be used.    
    - `BEDROCK_AGENT_TEST_UI_ICON` - (Optional) The favicon code, such as `:bar_chart:`. If this environment is not set, the generic icon in the above screenshot will be used.
    
4. Run the following command to start the Streamlit app:  
    `streamlit run app.py --server.port=8080 --server.address=localhost`
    

Once the app is started, you can access it in your web browser at `http://localhost:8080`.

As an example, here is the list of bash commands I run in bash inside VS Code to start the app for testing my forex rate agent (which you can [learn how to build](https://dev.to/aws-builders/building-a-basic-forex-rate-assistant-using-agents-for-amazon-bedrock-17fp) or [deploy using my Terraform configuration](https://dev.to/aws-builders/how-to-manage-an-amazon-bedrock-agent-using-terraform-1lag)):

```bash
cd amazon-bedrock-agent-test-ui
pip install -r requirements.txt
# Use a named profile created by the "aws configure sso" command
export AWS_PROFILE=AWSAdministratorAccess-<redacted>
export BEDROCK_AGENT_ID=WENOOVMMEK
export BEDROCK_AGENT_TEST_UI_TITLE="Forex Rate Assistant"
export BEDROCK_AGENT_TEST_UI_ICON=":currency_exchange:"
# Log in via the browser when prompted
aws sso login
streamlit run app.py --server.port=8080 --server.address=localhost
```

![The test UI](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/82dj9y1spkozopi04vug.png)

To stop the app, send an INT signal (Ctrl+C) in the prompt where you are running the `streamlit` command.

>⚠ On Windows, the `streamlit` command doesn't seem to end the process if you don't have the UI opened in the browser. If you run into this issue, simply go to `http://localhost:8080` in your browser, then hit Ctrl+C again in the prompt.

## Next steps

While the Streamlit app serves my purpose as it is, there are a few missing features which I will continue to add over time:

* Support for the use of [Knowledge Bases for Amazon Bedrock](https://docs.aws.amazon.com/bedrock/latest/userguide/knowledge-base.html) in an agent, such as displaying citations
    
* Support for [returning control to the agent developer](https://docs.aws.amazon.com/bedrock/latest/userguide/agents-returncontrol.html)
    
* Ability to switch between agents and aliases within the app
    

I also have not shown how to build and deploy the Streamlit app as a container in AWS, which I will perhaps demonstrate in another future blog post.

You are encountered to fork or copy the repository and built upon the existing code to suit your needs.

## Summary

In this blog post, I provided an introduction to a generic Streamlit UI that I built to facilitate more efficient testing of agents built with the Agents for Amazon Bedrock service. You can clone the repository and follow the instructions to run it locally, and improve upon the baseline as you see fit.

I will be adding more features and fixing bugs over time, so be sure to check out the repository from time to time. Be sure to follow the [Avangards Blog](https://blog.avangards.io) as I continue my journey with building generative AI applications using Amazon Bedrock.

Thanks for checking in!