Changed the way it handles history

app.py

@@ -1,5 +1,5 @@
-import os
 import streamlit as st
+import os
 from openai import OpenAI
 
 api_key = os.getenv("NVIDIANIM_API_KEY")
@@ -33,25 +33,22 @@ def get_llama_response(question):
     st.session_state.messages.append({"role": "assistant", "content": response_text})
     return response_text
 
-def generate_response():
-    question = st.session_state.user_input
-    st.session_state.user_input = ""
-    response = get_llama_response(question)
-    return response
+st.title("Ask Llama 3.1 405B on Nvidia NIM")
 
-# Display chat history
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
 
-# User input and response
-user_input = st.chat_input("Your message")
-if user_input:
-    st.session_state.user_input = user_input
-    with st.chat_message("user"):
-        st.markdown(user_input)
+if "user_input" not in st.session_state:
+    st.session_state.user_input = ""
+
+user_input = st.chat_input("Your message", key="user_input")
 
-    if st.button("Submit"):
-        response = generate_response()
+if st.button("Submit"):
+    if user_input:
+        st.session_state.messages.append({"role": "user", "content": user_input})
+        with st.chat_message("user"):
+            st.markdown(user_input)
+        response = get_llama_response(user_input)
         with st.chat_message("assistant"):
             st.markdown(response)