Fix stream response in multiple chatbots

pages/chatbot.py

@@ -34,21 +34,6 @@ def display_messages():
             with st.chat_message("System"):
                 st.write(message.content)
 
-def launchQuery(query: str = None):
-
-    # Initialize the assistant's response
-    full_response = st.write_stream(
-        st.session_state["assistant"].ask(
-            query,
-            prompt_system=st.session_state.prompt_system,
-            messages=st.session_state[chapter_session_key]["messages"] if "messages" in st.session_state[chapter_session_key] else [],
-            # variables=st.session_state["data_dict"],
-            # metadata=metadata
-        ))
-
-    # Temporary placeholder AI message in chat history
-    st.session_state[chapter_session_key]["messages"].append(AIMessage(content=full_response, kwargs={"model": st.session_state["assistant"].getReadableModel()}))
-    st.rerun()
 
 
 def show_prompts(prompts: list[str]):
@@ -59,6 +44,15 @@ def show_prompts(prompts: list[str]):
         if expander.button(p, key=f"button_{chapter_num}_{i}_{uuid.uuid4()}"):
             launchQuery(p)
 
+def inputSubmut():
+    user_query = st.session_state[f"input_{chapter_session_key}"]
+
+    
+    if user_query is not None and user_query != "":
+        st.session_state[chapter_session_key]["messages"].append(HumanMessage(content=user_query))
+        
+        # Stream and display response
+        launchQuery(user_query)
 
 def page():
 
@@ -113,15 +107,40 @@ def page():
     # Displaying messages
     display_messages()
 
+    chat_input_key = f"input_{chapter_session_key}"
 
-    user_query = st.chat_input("", key=f"input_{chapter_session_key}_{uuid.uuid4()}")
-    
-    if user_query is not None and user_query != "":
+    # Réinitialisation si la clé existe déjà (quand on change de page)
+    if chat_input_key in st.session_state:
+        del st.session_state[chat_input_key]
+
+
+    st.chat_input("", key=chat_input_key, on_submit=inputSubmut)
     
-        st.session_state[chapter_session_key]["messages"].append(HumanMessage(content=user_query))
-        
-        # Stream and display response
-        launchQuery(user_query)
+
+    global response_placeholder
+    response_placeholder = st.empty()
+
+
+    # Check if we need to rerun the app
+    if st.session_state.get('rerun', False):
+        st.session_state['rerun'] = False
+        st.rerun()
+
+
+def launchQuery(query: str = None):
+
+    # Initialize the assistant's response
+    full_response = response_placeholder.write_stream(
+            st.session_state["assistant"].ask(
+            query,
+            prompt_system=st.session_state.prompt_system,
+            messages=st.session_state[chapter_session_key]["messages"] if "messages" in st.session_state[chapter_session_key] else [],
+        )
+    )
+
+    # Temporary placeholder AI message in chat history
+    st.session_state[chapter_session_key]["messages"].append(AIMessage(content=full_response, kwargs={"model": st.session_state["assistant"].getReadableModel()}))
+    st.session_state['rerun'] = True
 
 
 page()