qchat_app.py

# created 2024-04-25 by Noah Coad ncoad@amazon.com
# demo video at https://youtu.be/9UkoTsXMnlQ

import streamlit as st, boto3

# parameters
aws_region = 'us-west-2'
q_app_id = 'be05d006-3c25-439e-b25d-781ad5b2xxxx'
q_user_id = 'AmazonQ-Administrator'

# setup app
amazon_q = boto3.client('qbusiness', aws_region)
st.set_page_config(page_title="Amazon Q Chatbot") #HTML title
st.title("Amazon Q Chatbot") #page title

# get user input
input_text = st.chat_input("Chat with your bot here")

# keep chat history
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

# populate chat history in UI
for message in st.session_state.chat_history:
    with st.chat_message(message["role"]):
        st.markdown(message["text"])

# when there's user input
if input_text:
    
    # show user msg
    with st.chat_message("user"): 
        st.markdown(input_text)

    with st.spinner("Thinking..."):
        # call out to Amazon Q
        kwargs = dict(applicationId=q_app_id,
        	userMessage=input_text,
        	userId = q_user_id,
        	conversationId = st.session_state.get('conversation_id'),
        	parentMessageId = st.session_state.get('parentmessage_id'))
        answer = amazon_q.chat_sync(**{k: v for k, v in kwargs.items() if v})
    
    # append user input to history
    st.session_state.chat_history.append({"role":"user", "text":input_text})

    # get Q response to user
    chat_response = answer['systemMessage'].replace("<answer> ", "").replace("</answer> ", "")
    st.session_state.conversation_id = answer['conversationId']
    st.session_state.parentmessage_id = answer['systemMessageId']
    
    # update message with citations
    if answer['sourceAttributions']:
        citations = {}
        for x in answer['sourceAttributions']:
            for segment in x['textMessageSegments']:
                if segment['endOffset'] not in citations:
                    citations[segment['endOffset']] = []
                citations[segment['endOffset']] += [x['citationNumber']]
        prev_offset = 0
        citations = sorted(citations.items(), key=lambda x: x[0])
        modified_message = ""
        for offset, nums in citations:
            modified_message += chat_response[prev_offset:offset] + " " + "".join([f"[{x}] " for x in nums])
            prev_offset = offset        

        # add references to end
        modified_message += "\n\n" + "".join([f"1. {x['url']}\n" for x in answer['sourceAttributions']])
    else: modified_message = chat_response

    # show Q response
    with st.chat_message("assistant"):
        st.markdown(modified_message)

    # append Q response to history
    st.session_state.chat_history.append({"role":"assistant", "text":modified_message})