Update app.py

app.py

@@ -7,7 +7,6 @@ from gtts import gTTS
 import google.generativeai as genai
 from io import BytesIO  # Import BytesIO
 
-
 # Set your API key
 api_key = "AIzaSyC70u1sN87IkoxOoIj4XCAPw97ae2LZwNM"  # Replace with your actual API key
 genai.configure(api_key=api_key)
@@ -81,6 +80,7 @@ def send_message():
     user_input = st.session_state.user_input
     uploaded_files = st.session_state.uploaded_files
     prompts = []
+    prompt_parts = []
 
     # Populate the prompts list with the existing chat history
     for entry in st.session_state['chat_history']:
@@ -88,38 +88,52 @@ def send_message():
             if 'text' in part:
                 prompts.append(part['text'])
             elif 'data' in part:
+                # Add the image in base64 format to prompt_parts for vision model
+                prompt_parts.append({"data": part['data'], "mime_type": "image/jpeg"})
                 prompts.append("[Image]")
 
     # Append the user input to the prompts list
     if user_input:
         prompts.append(user_input)
         st.session_state['chat_history'].append({"role": "user", "parts": [{"text": user_input}]})
+        # Also add the user text input to prompt_parts
+        prompt_parts.append({"text": user_input})
 
     # Handle uploaded files
     if uploaded_files:
         for uploaded_file in uploaded_files:
             base64_image = get_image_base64(Image.open(uploaded_file))
             prompts.append("[Image]")
+            prompt_parts.append({"data": base64_image, "mime_type": "image/jpeg"})
             st.session_state['chat_history'].append({
                 "role": "user",
                 "parts": [{"mime_type": uploaded_file.type, "data": base64_image}]
             })
 
+    # Determine if vision model should be used
+    use_vision_model = any(part.get('mime_type') == 'image/jpeg' for part in prompt_parts)
+
     # Set up the model and generate a response
-    model_name = 'gemini-pro-vision' if st.session_state.get('use_vision_model', False) else 'gemini-pro'
+    model_name = 'gemini-pro-vision' if use_vision_model else 'gemini-pro'
     model = genai.GenerativeModel(
         model_name=model_name,
         generation_config=generation_config,
         safety_settings=safety_settings
     )
     chat_history_str = "\n".join(prompts)
-    prompt_parts = [{"text": chat_history_str}]
-    response = model.generate_content([{"role": "user", "parts": prompt_parts}])
+    if use_vision_model:
+        # Include text and images for vision model
+        generated_prompt = {"role": "user", "parts": prompt_parts}
+    else:
+        # Include text only for standard model
+        generated_prompt = {"role": "user", "parts": [{"text": chat_history_str}]}
+
+    response = model.generate_content([generated_prompt])
     response_text = response.text if hasattr(response, "text") else "No response text found."
 
     # After generating the response from the model, append it to the chat history
     if response_text:
-        st.session_state['chat_history'].append({"role": "model", "parts": [{"text": response_text}]})
+        st.session_state['chat_history'].append({"role": "model", "parts":[{"text": response_text}]})
 
         # Convert the response text to speech
         tts = gTTS(text=response_text, lang='en')