Update app.py

app.py

@@ -32,7 +32,7 @@ def get_image_base64(image):
     img_str = base64.b64encode(buffered.getvalue()).decode()
     return img_str
 
-# Callback function to send message and clear input
+# Function to send message and clear input
 def send_message():
     user_input = st.session_state['user_input']
     uploaded_files = st.session_state['file_uploader']
@@ -40,39 +40,34 @@ def send_message():
     if user_input or uploaded_files:
         # Save user input to the chat history
         if user_input:
-            st.session_state['chat_history'].append(("User", user_input))
-
-        # Process text input
-        if user_input and not uploaded_files:
-            prompts = [{"role": "user", "parts": [{"text": user_input}]}]
-            model = genai.GenerativeModel(
-                model_name='gemini-pro',
-                generation_config=generation_config,
-                safety_settings=safety_settings
-            )
-            response = model.generate_content(prompts)
-            response_text = response.text if hasattr(response, "text") else "No response text found."
-            st.session_state['chat_history'].append(("Gemini", response_text))
-
-        # Process and save uploaded images to the chat history
+            st.session_state['chat_history'].append({"role": "user", "parts": [{"text": user_input}]})
+
+        # Process uploaded images
+        image_prompts = []
         if uploaded_files:
             for uploaded_file in uploaded_files:
                 image = Image.open(uploaded_file).convert("RGB")  # Ensure image is in RGB
                 image_base64 = get_image_base64(image)
-                st.session_state['chat_history'].append(("User", f"Uploaded image: {uploaded_file.name}"))
-
                 image_prompt = {"role": "user", "parts": [{"mime_type": uploaded_file.type, "data": image_base64}]}
-                model = genai.GenerativeModel(
-                    model_name='gemini-pro-vision',
-                    generation_config=generation_config,
-                    safety_settings=safety_settings
-                )
-                response = model.generate_content([image_prompt])
-                response_text = response.text if hasattr(response, "text") else "No response text found."
-                st.session_state['chat_history'].append(("Gemini", response_text))
+                image_prompts.append(image_prompt)
+            st.session_state['chat_history'].extend(image_prompts)
+
+        # Choose the appropriate model based on the input type
+        model_name = 'gemini-pro-vision' if uploaded_files else 'gemini-pro'
+        model = genai.GenerativeModel(
+            model_name=model_name,
+            generation_config=generation_config,
+            safety_settings=safety_settings
+        )
+
+        # Generate the response
+        response = model.generate_content(st.session_state['chat_history'])
+        response_text = response.text if hasattr(response, "text") else "No response text found."
+        st.session_state['chat_history'].append({"role": "model", "parts": [{"text": response_text}]})
 
         # Clear the input box after sending the message
         st.session_state.user_input = ""
+        st.session_state.file_uploader = None  # Clear uploaded files
 
 # Multiline text input for the user to send messages
 user_input = st.text_area("Enter your message here:", key="user_input", value="")
@@ -89,6 +84,10 @@ uploaded_files = st.file_uploader(
 send_button = st.button("Send", on_click=send_message)
 
 # Display the chat history
-for message in st.session_state['chat_history']:
-    role, text = message
-    st.markdown(f"**{role.title()}**: {text}")
+for entry in st.session_state['chat_history']:
+    role = entry["role"]
+    parts = entry["parts"][0]
+    if 'text' in parts:
+        st.markdown(f"**{role.title()}**: {parts['text']}")
+    elif 'data' in parts:
+        st.markdown(f"**{role.title()}**: (Image)")