Upload 2 files

app.py

@@ -1,6 +1,7 @@
 import cohere
 import streamlit as st
 from streamlit.components.v1 import html
+from streamlit_extras.stylable_container import stylable_container
 import re
 import urllib.parse
 
@@ -26,6 +27,21 @@ if "messages" not in st.session_state:
     st.session_state.messages = []
 
 def get_ai_response(prompt, chat_history):
+    st.session_state.is_streaming = True
+
+    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("ai"):
         penalty_kwargs = {
             "frequency_penalty" if penalty_type == "Frequency Penalty" else "presence_penalty": penalty_value
@@ -41,16 +57,17 @@ def get_ai_response(prompt, chat_history):
             p=p,
             **penalty_kwargs
         )
-        
-        response = ""
+
+        st.session_state.response = ""
         placeholder = st.empty()
         for event in stream:
             if event.event_type == "text-generation":
                 content = event.text
-                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):
@@ -95,7 +112,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]["text"])
+                new_content = st.text_area("Edit message", height=200, value=st.session_state.messages[i]["text"])
                 col1, col2 = st.columns([1, 1])
                 with col1:
                     if st.form_submit_button("Save"):
@@ -106,16 +123,6 @@ def display_messages():
                     if st.form_submit_button("Cancel"):
                         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:  # メッセージリストが空でないことを確認
-            prompt = st.session_state.messages[-1]["text"]
-            response = get_ai_response(prompt, st.session_state.messages[:-1])
-            st.session_state.messages.append({"role": "CHATBOT", "text": response})
-            st.session_state.retry_flag = False
-            st.rerun()
-        else:
-            st.session_state.retry_flag = False  # retry_flagをFalseに設定して処理を続行
 
 # Add sidebar for advanced settings
 with st.sidebar:
@@ -130,7 +137,7 @@ with st.sidebar:
             log_text += message["text"] + "\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"""
@@ -149,7 +156,7 @@ with st.sidebar:
     
     st.header("Advanced Settings")
     model = st.selectbox("Model", options=["command-r-plus", "command-r"], index=0)
-    preamble = st.text_area("Preamble", height=100)
+    preamble = st.text_area("Preamble", height=200)
     temperature = st.slider("Temperature", min_value=0.0, max_value=1.0, value=0.3, step=0.1)
     k = st.slider("Top-K", min_value=0, max_value=500, value=0, step=1)
     p = st.slider("Top-P", min_value=0.01, max_value=0.99, value=0.75, step=0.01)
@@ -190,16 +197,35 @@ with st.sidebar:
         else:
             st.warning("Please enter a valid API Key.")
 
+# After Stop generating
+if "is_streaming" in st.session_state and st.session_state.is_streaming:
+    st.session_state.messages.append({"role": "CHATBOT", "text": st.session_state.response})
+    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:
+        prompt = st.session_state.messages[-1]["text"]
+        response = get_ai_response(prompt, st.session_state.messages[:-1])
+        st.session_state.messages.append({"role": "CHATBOT", "text": response})
+        st.session_state.retry_flag = False
+        st.rerun()
+    else:
+        st.session_state.retry_flag = False
+
 if prompt := st.chat_input("What is up?"):
     chat_history = st.session_state.messages.copy()
-
+    
     with st.chat_message("user"):
         st.write(prompt)
-
-    response = get_ai_response(prompt, chat_history)
-        
     st.session_state.messages.append({"role": "USER", "text": prompt})
+    
+    response = get_ai_response(prompt, chat_history)
+    
     st.session_state.messages.append({"role": "CHATBOT", "text": response})
     st.rerun()

requirements.txt

@@ -1 +1,2 @@
-cohere
+cohere
+streamlit-extras