Integrating buttons and appended messages into st.chat_input()

Summary

Hello,

This is a continuation of a post I made earlier this week about adding buttons to st.chat_input() (Using buttons in st.chat_input()). I was successfully able to implement callback functions. I am now dealing with an issue of appended message session_states accordingly. For example, when I click on a button I want the function to run which will append a message. However, if a button is not clicked I still want the message to append. Right now it is appending two messages if I click the button, no messages appended if I don’t click but submit a new entry to the chat on the first time and one message if I don’t click button but submit new entry after the first time.

Steps to reproduce

Code snippet:

from funcs import send_feedback

def session_state_feedback(prompt, content, is_correct_answer, is_correct_document, document):
    st.session_state.button_clicked = True
    message = {
        "role": "assistant",
        "question": prompt,
        "content": content,
        "document": document,
        "is_correct_document": is_correct_document,
        "is_correct_answer": is_correct_answer
    }
    st.session_state.messages.append(message)

    send_feedback(
        query=prompt,
        answer_obj=dict(answer=content),
        is_correct_answer=is_correct_answer,
        is_correct_document=is_correct_document,
        document=document
    )

    st.session_state.button_clicked = True

# Initialize button_clicked in session state if it doesn't exist
if 'button_clicked' not in st.session_state:
    st.session_state.button_clicked = False

if prompt := st.chat_input():
    if st.session_state.message[-1]['role'] != 'assistant':
        with st.chat_message("assistant"):
            try:
                with st.spinner("Thinking..."):
                    button_col1, button_col2, button_col3, _ = st.columns([1,1,1,6])
                    button_col1.button("A", key=f"{documents[0]}1", help="Correct answer", on_click=session_state_feedback, args=[prompt, st.session_state.results[-1].get("answer"), True, True, documents[0]])
                    button_col2.button("B", key=f"{documents[0]}2", help="Wrong answer and wrong document", on_click=session_state_feedback, args=[prompt, st.session_state.results[-1].get("answer"), False, False, documents[0]])
                    button_col3.button("C", key=f"{documents[0]}3", help="Wrong answer, but right document", on_click=session_state_feedback, args=[prompt, st.session_state.results[-1].get("answer"), False, True, documents[0]])

                    # Only append message if no button was clicked during THIS rerun
                    if not st.session_state.button_clicked:
                        message = {
                            "role": "assistant",
                            "question": prompt,
                            "content": st.session_state.results[-1].get("answer"),
                            "document": documents[0],
                            "is_correct_document": is_correct_document, 
                            "is_correct_answer": is_correct_answer
                        }
                        st.session_state.messages.append(message)
            finally:
                # Reset the button_clicked flag for subsequent interactions
                st.session_state.button_clicked = False

Expected behavior:

A single message is appended to st.session_state.messages whether a button is clicked or if a button is not clicked but a new entry is added to chat input.

Actual behavior:

Two messages are appended if the button is clicked. Either no messages or one message is appended if a button is not clicked but a new entry is inputted.

Thank you for the help!

This topic was automatically closed 180 days after the last reply. New replies are no longer allowed.