Update app.py and README.md

README.md

@@ -1,7 +1,7 @@
 ---
-title: Cohere Chat Ui
-emoji: 🌍
-colorFrom: blue
+title: Cohere Chat Ui Dev
+emoji: 💻
+colorFrom: purple
 colorTo: gray
 sdk: streamlit
 sdk_version: 1.33.0
@@ -15,7 +15,7 @@ User Guide for Cohere Chat UI
 This application provides an interactive chat interface using Cohere's chat models. The following is an explanation of the main features and how to use the application.
 
 1. Entering the API Key
-   - Enter your Cohere API key in the text box located in the sidebar.
+   - Enter your Cohere API key in the text box.
    - You can obtain your API key from here: https://dashboard.cohere.com/api-keys
    - Once the API key is entered correctly, the application will become available for use.
 
@@ -23,10 +23,11 @@ This application provides an interactive chat interface using Cohere's chat mode
    - Type your message in the chat input box and submit it to receive a response from the Cohere model.
    - The chat history is displayed on the page, with user and Cohere model messages alternating.
 
-3. Editing and Deleting Messages
-   - Below each message, there are "Edit" and "Delete" buttons.
-   - Clicking the "Edit" button displays a text area where you can edit the content of the message. After making changes, click the "Save" button to save your edits.
+3. Editing, Deleting, and Retrying Messages
+   - Below each message, there are "Edit", "Delete", and "Retry" buttons.
+   - Clicking the "Edit" button displays a text area where you can edit the content of the message. After making changes, click the "Save" button to save your edits, or click the "Cancel" button to discard the changes.
    - Clicking the "Delete" button removes the corresponding message from the chat history.
+   - Clicking the "Retry" button, which appears only for the last message from the AI, allows you to regenerate the AI's response. This is useful if you want to get a different or updated response without modifying the input message.
 
 4. Copying Messages
    - Below each message, there is a "Copy" button.
@@ -54,7 +55,7 @@ Cohere Chat UIの使用説明書
 このアプリケーションは、Cohereのチャットモデルを使用して対話型のチャットインターフェースを提供します。以下は、アプリケーションの主要な機能と使用方法の説明です。
 
 1. API キーの入力
-   - サイドバーにあるテキストボックスにCohere APIキーを入力してください。
+   - テキストボックスにCohere APIキーを入力してください。
    - ここからAPIキーを入手してください: https://dashboard.cohere.com/api-keys
    - APIキーが正しく入力されると、アプリケーションが使用可能になります。
 
@@ -62,10 +63,11 @@ Cohere Chat UIの使用説明書
    - チャット入力ボックスにメッセージを入力し、送信するとCohereモデルからの応答が表示されます。
    - チャットの履歴はページ上に表示され、ユーザーとCohereモデルのメッセージが交互に表示されます。
 
-3. メッセージの編集と削除
-   - 各メッセージの下には、「Edit」と「Delete」のボタンがあります。
-   - 「Edit」ボタンをクリックすると、メッセージの内容を編集できるテキストエリアが表示されます。変更を加えた後、「Save」ボタンをクリックして編集を保存します。
+3. メッセージの編集、削除、リトライ
+   - 各メッセージの下には、「Edit」、「Delete」、「Retry」のボタンがあります。
+   - 「Edit」ボタンをクリックすると、メッセージの内容を編集できるテキストエリアが表示されます。変更を加えた後、「Save」ボタンをクリックして編集を保存するか、「Cancel」ボタンをクリックして変更を破棄します。
    - 「Delete」ボタンをクリックすると、対応するメッセージがチャットの履歴から削除されます。
+   - 「Retry」ボタンは、AIからの最後のメッセージにのみ表示されます。このボタンをクリックすると、AIの応答を再生成できます。これは、入力メッセージを変更せずに、別の応答や更新された応答を取得したい場合に便利です。
 
 4. メッセージのコピー
    - 各メッセージの下には、「Copy」ボタンがあります。

app.py

@@ -4,22 +4,51 @@ import re
 
 st.title("Cohere Chat UI")
 
-api_key = st.sidebar.text_input("API Key", type="password")
-if api_key:
-    client = cohere.Client(api_key=api_key)
+if "api_key" not in st.session_state:
+    api_key = st.text_input("Enter your API Key", type="password")
+    if api_key:
+        st.session_state.api_key = api_key
+        client = cohere.Client(api_key=api_key)
+        st.rerun()
+    else:
+        st.warning("Please enter your API key to use the app. You can obtain your API key from here: https://dashboard.cohere.com/api-keys")
+        st.stop()
 else:
-    st.warning("Please enter your API key to use the app. You can obtain your API key from here: https://dashboard.cohere.com/api-keys")
-    st.stop()
+    client = cohere.Client(api_key=st.session_state.api_key)
 
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
+def get_ai_response(prompt, chat_history):
+    with st.chat_message("ai"):
+        penalty_kwargs = {
+            "frequency_penalty" if penalty_type == "Frequency Penalty" else "presence_penalty": penalty_value
+        }
+        stream = client.chat_stream(
+            message=prompt,
+            model=model,
+            preamble=preamble,
+            chat_history=chat_history,
+            temperature=temperature,
+            **penalty_kwargs
+        )
+        
+        response = ""
+        placeholder = st.empty()
+        for event in stream:
+            if event.event_type == "text-generation":
+                content = event.text
+                response += content
+                placeholder.markdown(response)
+    
+    return response
+
 def display_messages():
     for i, message in enumerate(st.session_state.messages):
         name = "user" if message["role"] == "USER" else "ai"
         with st.chat_message(name):
             st.markdown(message["text"])
-            col1, col2, col3 = st.columns([1, 1, 1])
+            col1, col2, col3, col4 = st.columns([1, 1, 1, 1])
             with col1:
                 if st.button("Edit", key=f"edit_{i}_{len(st.session_state.messages)}"):
                     st.session_state.edit_index = i
@@ -37,14 +66,33 @@ def display_messages():
                     copy_text = message["text"]
                     st.code(copy_text, language='plain')
 
+            if i == len(st.session_state.messages) - 1 and message["role"] == "CHATBOT":
+                with col4:
+                    if st.button("Retry", key=f"retry_{i}_{len(st.session_state.messages)}"):
+                        del st.session_state.messages[-1]
+                        st.session_state.retry_flag = True
+                        st.rerun()
+        
         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"])
-                if st.form_submit_button("Save"):
-                    st.session_state.messages[i]["text"] = new_content
-                    del st.session_state.edit_index
-                    st.rerun()
-display_messages()
+                col1, col2 = st.columns([1, 1])
+                with col1:
+                    if st.form_submit_button("Save"):
+                        st.session_state.messages[i]["text"] = new_content
+                        del st.session_state.edit_index
+                        st.rerun()
+                with col2:
+                    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:
+        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()
 
 # Add sidebar for advanced settings
 with st.sidebar:
@@ -92,32 +140,25 @@ with st.sidebar:
             st.session_state.messages.append({"role": role, "text": text.strip()})
         st.rerun()
 
+    st.header("Change API Key")
+    new_api_key = st.text_input("Enter new API Key", type="password")
+    if st.button("Update API Key"):
+        if new_api_key:
+            st.session_state.api_key = new_api_key
+            client = cohere.Client(api_key=new_api_key)
+            st.success("API Key updated successfully!")
+        else:
+            st.warning("Please enter a valid API Key.")
+
+display_messages()
+
 if prompt := st.chat_input("What is up?"):
     chat_history = st.session_state.messages.copy()
 
     with st.chat_message("user"):
         st.write(prompt)
 
-    with st.chat_message("ai"):
-        penalty_kwargs = {
-            "frequency_penalty" if penalty_type == "Frequency Penalty" else "presence_penalty": penalty_value
-        }
-        stream = client.chat_stream(
-            message=prompt,
-            model=model,
-            preamble=preamble,
-            chat_history=chat_history,
-            temperature=temperature,
-            **penalty_kwargs
-        )
-        
-        response = ""
-        placeholder = st.empty()
-        for event in stream:
-            if event.event_type == "text-generation":
-                content = event.text
-                response += content
-                placeholder.markdown(response)
+    response = get_ai_response(prompt, chat_history)
         
     st.session_state.messages.append({"role": "USER", "text": prompt})
     st.session_state.messages.append({"role": "CHATBOT", "text": response})