highlight and rate text in LLM response within a chat app

Hi!

For my Streamlit-based LLM chat app, I’d like to have a feature that allows the user to select and rate certain pieces of text from the LLM’s response. The result should be similar to the screenshots below:


The corresponding code can be found below:

from streamlit_annotation_tools import text_highlighter
import streamlit as st

text = "Yesterday, at 3 PM, Emily Johnson and Michael Smith met at the Central Park in New York to  discuss the merger between TechCorp and Global Solutions. The deal, worth approximately 500  million dollars, is expected to significantly impact the tech industry. Later, at 6 PM, they joined a  conference call with the CEO of TechCorp, David Brown, who was in London for a technology  summit. During the call, they discussed the market trends in Asia and Europe and planned for the  next quarterly meeting, which is scheduled for January 15th, 2024, in Paris."

highlights = [[{"start":0,"end":0, "label":"SELECTION"}]]
highlights = text_highlighter(text, highlights)

st.write("Highlighted Text:")
st.write(highlights)

key_count = 0
feedback_list = []
feedback_text_list = []

with st.popover("Open popover", use_container_width=True):
    if highlights:
        for item in highlights:
            for el in item:
                st.write(f"{el['label']}")
                feedback = st.feedback(key=key_count)
                feedback_text = st.text_input(
                    "Feedback for highlighted text",
                    key=f"feedback_text_{key_count}",
                    placeholder="Enter your feedback here..."
                )
                key_count += 1 
                feedback_list.append(feedback)
                feedback_text_list.append(feedback_text)


st.write("Feedback:")
for i, feedback in enumerate(feedback_list):
    st.write(f"Feedback {i + 1}: {feedback}")
    if feedback_text_list[i]:
        st.write(f"Feedback Text {i + 1}: {feedback_text_list[i]}")
    else:
        st.write(f"Feedback Text {i + 1}: No text provided")

However, when I start integrating the text highlight feature within a chat app as follows, the highlighting feature is only shown very quickly and then disappears. Therefore, the user cannot really use it. Any ideas how to integrate this feature properly into a chat app?

from streamlit_annotation_tools import text_highlighter
import streamlit as st
import time
import random

# Streamed response emulator
def response_generator():
    response = random.choice(
        [
            "Hello there! How can I assist you today?",
            "Hi, human! Is there anything I can help you with?",
            "Do you need help?",
        ]
    )
    for word in response.split():
        yield word + " "
        time.sleep(0.05)


st.title("Simple chat")

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])


# Accept user input
if prompt := st.chat_input("What is up?"):
    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": prompt})
    # Display user message in chat message container
    with st.chat_message("user"):
        st.markdown(prompt)

    # Display assistant response in chat message container
    with st.chat_message("assistant"):
        response = st.write_stream(response_generator())
        highlights = text_highlighter(response)
    

    # Add assistant response to chat history
    st.session_state.messages.append({"role": "assistant", "content": response})

I have also looked at this other thread which tries to achieve a similar thing but it has other issues with the st.dialog and the used highlighting library.