Changed app.py to preserve conversation history

app.py

@@ -1,6 +1,6 @@
-from openai import OpenAI
-import streamlit as st
 import os
+import streamlit as st
+from openai import OpenAI
 
 api_key = os.getenv("NVIDIANIM_API_KEY")
 
@@ -11,32 +11,47 @@ client = OpenAI(
 
 model_name = "meta/llama-3.1-405b-instruct"
 
-def get_llama_response(question):
-  completion = client.chat.completions.create(
-      model=model_name,
-      messages=[{"role": "user", "content": question}],
-      temperature=0.2,
-      top_p=0.7,
-      max_tokens=1024,
-      stream=True
-  )
-  
-  response = ""
-  for chunk in completion:
-    if chunk.choices[0].delta.content is not None:
-      response += chunk.choices[0].delta.content
-  return response.strip()
-
-st.title("Ask Llama 3.1 405B on Nvidia NIM")
-user_question = st.text_input("Enter your question:")
-
-if st.button("Submit"):
-  if user_question:
-    llama_response = get_llama_response(user_question)
-    st.write("**Llama 3.1 405B Response:**")
-    st.write(llama_response)
-  else:
-    st.warning("Please enter a question.")
-
-
+if "messages" not in st.session_state:
+    st.session_state.messages = []
 
+def get_llama_response(question):
+    st.session_state.messages.append({"role": "user", "content": question})
+    response = client.chat.completions.create(
+        model=model_name,
+        messages=st.session_state.messages,
+        temperature=0.2,
+        top_p=0.7,
+        max_tokens=1024,
+        stream=True
+    )
+
+    response_text = ""
+    for chunk in response:
+        if chunk.choices[0].delta.content is not None:
+            response_text += chunk.choices[0].delta.content
+
+    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
+
+# 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 st.button("Submit"):
+        response = generate_response()
+        with st.chat_message("assistant"):
+            st.markdown(response)