HTAP Summit 2024 session replays are now live!Access Session Replays
Sign In Start for Free
Start for Free

Creating a Retrieval-Augmented Generation (RAG) based chatbot using AWS Bedrock, Streamlit, and TiDB Serverless can significantly enhance the user experience by providing efficient and accurate responses. This tutorial will guide you through the process of building this robust chatbot step-by-step.

Introduction

Combining AWS Bedrock for model management, Streamlit for the frontend interface, and TiDB Serverless for vector storage creates a powerful and scalable RAG-based chatbot. This tutorial assumes you have basic knowledge of Python and database management.

Prerequisites

Before you begin, ensure you have the following:

  • A TiDB Serverless cluster with vector search enabled.
  • Python 3.8 or higher.
  • An AWS account with access to Amazon Bedrock.

Step 1: Setting Up Your Environment

First, clone the GitHub repository containing the necessary files for this project:

git clone https://github.com/Yoshiitaka/TiDB-Bedrock.git

Next, create and activate a virtual environment:

python -m venv .venv
source .venv/bin/activate

Install the required libraries and packages:

pip install -r requirements.txt

Step 2: Configure Environment Variables

Create a .env file in the root directory with your TiDB and AWS credentials. Use the provided sample-env file as a reference:

TIDB_HOSTNAME='your_tidb_hostname'
TIDB_USERNAME='your_tidb_username'
TIDB_PASSWORD='your_tidb_password'
TIDB_DATABASE_NAME='your_tidb_database_name'

Step 3: Preparing the Database

Before running the chatbot, you need to prepare the TiDB database by embedding vector data. This step ensures the chatbot can retrieve relevant information efficiently.

Run the preparation script:

python prepare.py

Step 4: Developing the Chatbot

Here’s the complete code to develop the RAG-based chatbot:

import logging
import sys
import streamlit as st
from sqlalchemy import URL
from llama_index.core import VectorStoreIndex
from llama_index.core.settings import Settings
from llama_index.llms.bedrock import Bedrock
from llama_index.vector_stores.tidbvector import TiDBVectorStore
from llama_index.embeddings.bedrock import BedrockEmbedding
import os
from dotenv import load_dotenv
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logger = logging.getLogger()

# Load environment variables
load_dotenv()
tidb_username = os.environ['TIDB_USERNAME']
tidb_password = os.environ['TIDB_PASSWORD']
tidb_host = os.environ['TIDB_HOSTNAME']
tidb_database = os.environ['TIDB_DATABASE_NAME']

# Initialize LlamaIndex with Amazon Bedrock
llm = Bedrock(model="anthropic.claude-3-sonnet-20240229-v1:0")
embed_model = BedrockEmbedding(model="amazon.titan-embed-text-v1")

Settings.llm = llm
Settings.embed_model = embed_model

# Initialize TiDB Vector Store
if 'tidb_vec_index' not in st.session_state:
    tidb_connection_url = URL(
        "mysql+pymysql",
        username=tidb_username,
        password=tidb_password,
        host=tidb_host,
        port=4000,
        database=tidb_database,
        query={"ssl_verify_cert": True, "ssl_verify_identity": True},
    )

    tidbvec = TiDBVectorStore(
        connection_string=tidb_connection_url,
        table_name="llama_index_rag",
        distance_strategy="cosine",
        vector_dimension=1536,
        drop_existing_table=False,
    )

    tidb_vec_index = VectorStoreIndex.from_vector_store(tidbvec)
    st.session_state['tidb_vec_index'] = tidb_vec_index

query_engine = st.session_state['tidb_vec_index'].as_query_engine(streaming=True)

# Streamlit UI Configuration
st.set_page_config(page_title='TiDB & Bedrock DEMO')

def clear_screen():
    st.session_state.messages = [{"role": "assistant", "content": "I'm learning about TiDB. Feel free to ask me anything!"}]

with st.sidebar:
    st.title('RAG Application with TiDB and Bedrock 🤖')
    st.divider()
    st.image('public/tidb-logo-with-text.png', caption='TiDB')
    st.image('public/bedrock.png', caption='Amazon Bedrock')
    st.button('Clear Screen', on_click=clear_screen)

if "messages" not in st.session_state.keys():
    st.session_state.messages = [{"role": "assistant", "content": "I'm learning about TiDB. Feel free to ask me anything!"}]

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])

if prompt := st.chat_input():
    with st.chat_message("user"):
        st.markdown(prompt)
        st.session_state.messages.append({"role": "user", "content": prompt})

    with st.chat_message("assistant"):
        placeholder = st.empty()
        full_response = ''
        streaming_response = query_engine.query(prompt)
        for chunk in streaming_response.response_gen:
            full_response += chunk
            placeholder.markdown(full_response)
        placeholder.markdown(full_response)
        st.session_state.messages.append({"role": "assistant", "content": full_response})

Step 5: Running the Application

Start the Streamlit server to interact with the chatbot:

streamlit run main.py

Open your browser and navigate to http://localhost:8501/ to start using the chatbot.

Conclusion

By following this tutorial, you have successfully built a RAG-based chatbot leveraging AWS Bedrock, Streamlit, and TiDB Serverless. This powerful combination provides an efficient and scalable solution for various applications. Explore further by integrating additional features and fine-tuning the models to fit your specific needs.

Feel free to experiment and expand this setup to suit your requirements. Happy coding!

You can also follow tutorials to learn how to build AI apps and how to use TiDB as vector store:


Last updated June 23, 2024

Spin up a Serverless database with 25GiB free resources.

Start Right Away