Update app.py

app.py

@@ -33,12 +33,11 @@ safety_settings = [
 ]
 
 # Initialize the chat history and the model instance using session state
-if "chat_history" not in st.session_state:
-  st.session_state.chat_history = []
-if "model" not in st.session_state:
-  st.session_state.model = genai.GenerativeModel(model_name="gemini-pro",
+# Use the get function to avoid KeyError if the key does not exist
+chat_history = st.session_state.get("chat_history", [])
+model = st.session_state.get("model", genai.GenerativeModel(model_name="gemini-pro",
                               generation_config=generation_config,
-                              safety_settings=safety_settings)
+                              safety_settings=safety_settings))
 
 # Create a streamlit interface with a title, a description, and a chat container
 st.title("Gemini API") # the title of the interface
@@ -46,7 +45,7 @@ st.markdown("A simple interface for Gemini API that explains code snippets in na
 chat_container = st.empty() # the chat container to display chat messages
 
 # Display chat messages from history on app rerun
-for message in st.session_state.chat_history:
+for message in chat_history:
   role = message["role"]
   content = message["content"]
   with chat_container.chat_message(role):
@@ -58,18 +57,19 @@ user_input, attachment = st.chat_input("Enter a message or a code snippet")
 if user_input or attachment:
   # Add user message to chat history
   if user_input:
-    st.session_state.chat_history.append({"role": "user", "content": user_input})
+    chat_history.append({"role": "user", "content": user_input})
   # Display user message in chat container
   with chat_container.chat_message("user"):
     st.markdown(user_input)
   if attachment:
     # You can process the attachment here or save it as needed
     # For example, you can save it to a temporary folder and store the file path in the chat history
+    # Use the attachment.name attribute to get the file name
     attachment_path = f"attachments/{attachment.name}"
     with open(attachment_path, "wb") as f:
       f.write(attachment.read())
     # Add attachment message to chat history
-    st.session_state.chat_history.append({"role": "user", "content": f"Attachment: {attachment.name}"})
+    chat_history.append({"role": "user", "content": f"Attachment: {attachment.name}"})
     # Display attachment message in chat container
     with chat_container.chat_message("user"):
       st.markdown(f"Attachment: {attachment.name}")
@@ -78,7 +78,7 @@ if user_input or attachment:
   with chat_container.chat_message("assistant"):
     with st.spinner("Generating a response..."):
       # Start a chat session with the model
-      convo = st.session_state.model.start_chat(history=[])
+      convo = model.start_chat(history=[])
       # Send the user input to the model
       convo.send_message(user_input)
       # Get the response from the model
@@ -86,4 +86,4 @@ if user_input or attachment:
       # Display the response in the chat container
       st.markdown(response)
       # Add assistant message to chat history
-      st.session_state.chat_history.append({"role": "assistant", "content": response})
+      chat_history.append({"role": "assistant", "content": response})