Update app.py

app.py

@@ -4,6 +4,7 @@ from streamlit.components.v1 import html
 from streamlit_extras.stylable_container import stylable_container
 import re
 import urllib.parse
+import traceback
 
 st.title("Nvidia Chat UI")
 
@@ -19,7 +20,7 @@ if "api_key" not in st.session_state:
             st.warning("Please enter your API key correctly.")
             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.warning("Please enter your API key to use the app. You can obtain your API key from here: https://build.nvidia.com/nvidia/nemotron-4-340b-instruct")
         st.stop()
 else:
     client = openai.OpenAI(base_url = "https://integrate.api.nvidia.com/v1",
@@ -34,44 +35,48 @@ def get_ai_response(prompt, chat_history):
 
     chat_history.insert(0, {"role": "system", "content": system_prompt})
     chat_history.append({"role": "user", "content": prompt})
-    
-    with st.chat_message("assistant", avatar=st.session_state.assistant_avatar):
-        stream = client.chat.completions.create(
-            model="nvidia/nemotron-4-340b-instruct",
-            messages=chat_history,
-            temperature=temperature,
-            top_p=top_p,
-            max_tokens=1024,
-            stream=True,
-        )
-
-        placeholder = st.empty()
-
-        with stylable_container(
-            key="stop_generating",
-            css_styles="""
-                button {
-                    position: fixed;
-                    bottom: 100px;
-                    left: 50%;
-                    transform: translateX(-50%);
-                    z-index: 1;
-                }
-                """,
-        ):
-            st.button("Stop generating")
 
-        shown_message = ""
-        
-        for chunk in stream:
-            if chunk.choices[0].delta.content is not None:
-                content = chunk.choices[0].delta.content
-                st.session_state.response += content
-                shown_message += content.replace("\n", "  \n")\
-                                        .replace("<", "\\<")\
-                                        .replace(">", "\\>")
-                placeholder.markdown(shown_message)
+    try:
+        with st.chat_message("assistant", avatar=st.session_state.assistant_avatar):
+            stream = client.chat.completions.create(
+                model=model,
+                messages=chat_history,
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=1024,
+                stream=True,
+            )
+    
+            placeholder = st.empty()
+    
+            with stylable_container(
+                key="stop_generating",
+                css_styles="""
+                    button {
+                        position: fixed;
+                        bottom: 100px;
+                        left: 50%;
+                        transform: translateX(-50%);
+                        z-index: 1;
+                    }
+                    """,
+            ):
+                st.button("Stop generating")
+    
+            shown_message = ""
+            
+            for chunk in stream:
+                if chunk.choices[0].delta.content is not None:
+                    content = chunk.choices[0].delta.content
+                    st.session_state.response += content
+                    shown_message += content.replace("\n", "  \n")\
+                                            .replace("<", "\\<")\
+                                            .replace(">", "\\>")
+                    placeholder.markdown(shown_message) 
 
+    except:
+        st.session_state.response +=  "\n!Exception\n" + traceback.format_exc()
+    
     st.session_state.is_streaming = False
     return st.session_state.response
 
@@ -199,7 +204,7 @@ with st.sidebar:
         st.session_state.is_delete_mode = st.toggle("Enable Delete button")
         
         st.header("Advanced Settings")
-        model = st.selectbox("Model", options=["command-r-plus", "command-r"], index=0)
+        model = st.selectbox("Model", options=["nvidia/nemotron-4-340b-instruct"], index=0)
         system_prompt = st.text_area("System Prompt", height=200)
         temperature = st.slider("Temperature", min_value=0.0, max_value=1.0, value=0.3, step=0.1)
         top_p = st.slider("Top-P", min_value=0.01, max_value=1.00, value=1.00, step=0.01)