Update app.py

app.py

@@ -1,80 +1,87 @@
 import streamlit as st
 import google.generativeai as genai
 
-# Configure API key
-genai.configure(api_key="AIzaSyC70u1sN87IkoxOoIj4XCAPw97ae2LZwNM")  
+# API key 
+genai.configure(api_key="AIzaSyC70u1sN87IkoxOoIj4XCAPw97ae2LZwNM")
 
 # Model settings
 generation_config = {
   "temperature": 0.9,
   "top_p": 1,
-  "top_k": 1,  
-  "max_output_tokens": 2048   
+  "top_k": 1,
+  "max_output_tokens": 2048    
 }
 
 safety_settings = [
-  {
-     "category": "HARM_CATEGORY_HARASSMENT",
-     "threshold": "BLOCK_MEDIUM_AND_ABOVE"  
-  },
-  { 
-     "category": "HARM_CATEGORY_HATE_SPEECH",
-     "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-  },
-  {
-     "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", 
-     "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-  },
-  {
-     "category": "HARM_CATEGORY_DANGEROUS_CONTENT",  
-     "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-  }
+  {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+  {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+  {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+  {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"}
 ]
 
 model = genai.GenerativeModel(
-   model_name="gemini-pro",
-   generation_config=generation_config, 
-   safety_settings=safety_settings  
+  model_name="gemini-pro", 
+  generation_config=generation_config,
+  safety_settings=safety_settings
 )
 
 # Chatbot interface
-st.title("Gemini API Chatbot")   
+st.title("Gemini API Chatbot")
 
-chat_history = st.session_state.get("chat_history", [])
+chat_history = st.session_state.get("chat_history", []) 
 
-chat_container = st.container()  
+chat_container = st.container()
 
 for message in chat_history:
   # Display message
-  role = message["role"]
-  content = message["content"] 
+  if message.parts:
+    role = "assistant" if message.role == "system" else "user"
+    text = message.parts[0].text 
+  else:
+    role = message.role
+    text = message.content
   
   with chat_container.empty():
-    st.markdown(f"**{role}:** {content}")
+    st.markdown(f"**{role}:** {text}")
 
-# Get user input  
-user_input = st.text_input("Ask the assistant")   
+# Input  
+user_input = st.text_input("You")
 
 if user_input:
 
-  # Process input
-  chat_history.append({"role": "user", "content": user_input})
-   
-  # Display input  
+  # Create message
+  user_message = genai.Content(
+      parts=[genai.Part(text=user_input)], 
+      role=genai.Role.USER
+  )
+
+  # Append message  
+  chat_history.append(user_message)
+
+  # Display message
   with chat_container.empty():
     st.markdown(f"**user:** {user_input}")
 
   # Get response
   with st.spinner("Thinking..."):
-    convo = model.start_chat(history=chat_history)
-    response = convo.last.text
+    convo = model.start_chat(chat_history)
+    response = convo.last
 
-  # Process response 
-  chat_history.append({"role": "assistant", "content": response})
+  # Extract response text
+  response_text = response.parts[0].text
   
-  # Display response
+  # Create response message
+  assistant_message = genai.Content(
+      parts=[genai.Part(text=response_text)],
+      role=genai.Role.ASSISTANT
+  )
+
+  # Append response 
+  chat_history.append(assistant_message)
+
+  # Display response 
   with chat_container.empty():
-    st.markdown(f"**assistant:** {response}")
-    
+    st.markdown(f"**assistant:** {response_text}")
+  
   # Update session state
   st.session_state["chat_history"] = chat_history