Added the Stop generation button

app.py

@@ -1,6 +1,7 @@
 from premai import Prem
 import streamlit as st
 from streamlit.components.v1 import html
+from streamlit_extras.stylable_container import stylable_container
 import re
 import urllib.parse
 
@@ -28,6 +29,22 @@ if "messages" not in st.session_state:
     st.session_state.messages = []
 
 def get_ai_response(messages):
+    st.session_state.is_streaming = True
+    st.session_state.response = ""
+
+    with stylable_container(
+            key="stop_generating",
+            css_styles="""
+                button {
+                    position: fixed;
+                    bottom: 100px;
+                    left: 50%;
+                    transform: translateX(-50%);
+                }
+                """,
+        ):
+            st.button("Stop generating")
+
     with st.chat_message("assistant"):
         penalty_kwargs = {
             "frequency_penalty" if penalty_type == "Frequency Penalty" else "presence_penalty": penalty_value
@@ -44,16 +61,15 @@ def get_ai_response(messages):
             **penalty_kwargs,
         )
         
-        response = ""
-        
         placeholder = st.empty()
         for chunk in stream:
             if chunk.choices[0].delta["content"]:
                 content = chunk.choices[0].delta["content"]
-                response += content
-                placeholder.markdown(response)
-    
-    return response
+                st.session_state.response += content
+                placeholder.markdown(st.session_state.response)
+
+    st.session_state.is_streaming = False
+    return st.session_state.response
 
 def display_messages():
     for i, message in enumerate(st.session_state.messages):
@@ -97,7 +113,7 @@ def display_messages():
         
         if "edit_index" in st.session_state and st.session_state.edit_index == i:
             with st.form(key=f"edit_form_{i}_{len(st.session_state.messages)}"):
-                new_content = st.text_area("Edit message", value=st.session_state.messages[i]["content"])
+                new_content = st.text_area("Edit message", height=200, value=st.session_state.messages[i]["content"])
                 col1, col2 = st.columns([1, 1])
                 with col1:
                     if st.form_submit_button("Save"):
@@ -109,24 +125,8 @@ def display_messages():
                         del st.session_state.edit_index
                         st.rerun()
     
-    if "retry_flag" in st.session_state and st.session_state.retry_flag == True:
-        if len(st.session_state.messages) > 0:  # メッセージリストが空でないことを確認
-            response = get_ai_response(st.session_state.messages)
-            st.session_state.messages.append({"role": "assistant", "content": response})
-
-            if response:
-                st.session_state.is_error = False
-            else:
-                st.session_state.is_error = True
-            
-            st.session_state.retry_flag = False
-            st.rerun()
-                
-        else:
-            st.session_state.retry_flag = False  # retry_flagをFalseに設定して処理を続行
-    
     if "is_error" in st.session_state and st.session_state.is_error:
-        st.warning("""
+        st.error("""
                    Something went wrong. To resolve this error:
                    1. Use the Retry button.
                    2. Update your API key or Project ID correctly.
@@ -148,7 +148,7 @@ with st.sidebar:
             log_text += message["content"] + "\n\n"
     log_text = log_text.rstrip("\n")
         
-    # PythonでURLエンコード
+    # Encode the string to escape
     log_text_escaped = urllib.parse.quote(log_text)
 
     copy_log_button_html = f"""
@@ -208,7 +208,7 @@ with st.sidebar:
                   "remm-slerp-l2-13b",
                   ]
     model = st.selectbox("Model", options=model_list, index=0)
-    system_prompt = st.text_area("System prompt", height=100)
+    system_prompt = st.text_area("System prompt", height=200)
     temperature = st.slider("Temperature", min_value=0.0, max_value=1.0, value=1.0, step=0.1)
     top_p = st.slider("Top-P", min_value=0.01, max_value=0.99, value=0.75, step=0.01)
     penalty_type = st.selectbox("Penalty Type", options=["Frequency Penalty", "Presence Penalty"])
@@ -256,11 +256,37 @@ with st.sidebar:
         else:
             st.warning("Please enter a valid Project ID.")
 
+# After Stop generating
+if "is_streaming" in st.session_state and st.session_state.is_streaming:
+    st.session_state.messages.append({"role": "assistant", "content": st.session_state.response})
+    st.session_state.is_error = False
+    st.session_state.is_streaming = False
+    if "retry_flag" in st.session_state and st.session_state.retry_flag:
+        st.session_state.retry_flag = False
+    st.rerun()
+
 display_messages()
 
+# After Retry
+if "retry_flag" in st.session_state and st.session_state.retry_flag == True:
+    if len(st.session_state.messages) > 0:
+        response = get_ai_response(st.session_state.messages)
+        st.session_state.messages.append({"role": "assistant", "content": response})
+
+        if response:
+            st.session_state.is_error = False
+        else:
+            st.session_state.is_error = True
+
+        st.session_state.retry_flag = False
+        st.rerun()
+            
+    else:
+        st.session_state.retry_flag = False
+
 if prompt := st.chat_input("What is up?"):
+    st.session_state.messages.append({"role": "user", "content": prompt})
     messages = st.session_state.messages.copy()
-    messages.append({"role": "user", "content": prompt})
 
     with st.chat_message("user"):
         st.write(prompt)
@@ -272,7 +298,6 @@ if prompt := st.chat_input("What is up?"):
     else:
         st.session_state.is_error = True
 
-    st.session_state.messages.append({"role": "user", "content": prompt})
     st.session_state.messages.append({"role": "assistant", "content": response})
     
     st.rerun()