Created
April 26, 2024 02:00
-
-
Save noahcoad/11d175ebe3047e2fd2cf132f32203185 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# 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}) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment