app.py

app.py

@@ -74,12 +74,11 @@ def contract_chat_section(tab, assistant_id, session_key, input_key):
         for idx, message in enumerate(st.session_state[session_key]):
             role, content = message["role"], message["content"]
             if role == "assistant":
-                col1, col2 = st.columns([5, 1])
-                with col1:
+                with st.container():
                     st.chat_message(role).write(content)
-                with col2:
                     if st.button("🚩 Flag", key=f"flag_{session_key}_{idx}"):
-                        save_flagged_response(st.session_state[session_key][idx-1]["content"], content)
+                        user_query = st.session_state[session_key][idx-1]["content"] if idx > 0 else "Unknown"
+                        save_flagged_response(user_query, content)
             else:
                 st.chat_message(role).write(content)
 
@@ -109,8 +108,11 @@ def contract_chat_section(tab, assistant_id, session_key, input_key):
 
                 messages = client.beta.threads.messages.list(thread_id=thread_id)
                 assistant_message = messages.data[0].content[0].text.value
-                st.chat_message("assistant").write(assistant_message)
                 st.session_state[session_key].append({"role": "assistant", "content": assistant_message})
+                with st.container():
+                    st.chat_message("assistant").write(assistant_message)
+                    if st.button("🚩 Flag", key=f"flag_response_{session_key}_{len(st.session_state[session_key])}"):
+                        save_flagged_response(prompt, assistant_message)
             except Exception as e:
                 st.error(f"Error: {str(e)}")