Update app.py

app.py

@@ -23,6 +23,8 @@ if 'chat_history' not in st.session_state:
     st.session_state['chat_history'] = []
 if 'file_uploader_key' not in st.session_state:
     st.session_state['file_uploader_key'] = str(uuid.uuid4())
+if 'using_vision_model' not in st.session_state:
+    st.session_state['using_vision_model'] = False
 
 # UI layout
 st.title("Gemini Chatbot")
@@ -38,6 +40,7 @@ def get_image_base64(image):
 # Function to clear conversation
 def clear_conversation():
     st.session_state['chat_history'] = []
+    st.session_state['using_vision_model'] = False
     st.session_state['file_uploader_key'] = str(uuid.uuid4())
 
 # Function to send message and clear input
@@ -45,40 +48,35 @@ def send_message():
     user_input = st.session_state.user_input
     uploaded_files = st.session_state.uploaded_files
 
-    # Process text input for multi-turn conversation
-    if user_input:
-        st.session_state['chat_history'].append({"role": "user", "parts": [{"text": user_input}]})
-
-    # Check if images are uploaded
-    if uploaded_files:
-        # Prepare image prompts for single-turn conversation
-        image_prompts = [{
-            "role": "user",
-            "parts": [{"mime_type": uploaded_file.type, "data": get_image_base64(Image.open(uploaded_file))}]
-        } for uploaded_file in uploaded_files]
-        
+    # If images are uploaded or we have already used the vision model, 
+    # bundle the entire conversation history into a single-turn interaction.
+    if uploaded_files or st.session_state['using_vision_model']:
+        st.session_state['using_vision_model'] = True
+        prompts = []
+
+        # Add all previous chat history to the prompt
+        for entry in st.session_state['chat_history']:
+            parts = entry["parts"][0]
+            prompts.append({
+                "role": entry["role"],
+                "parts": [{"text": parts.get("text"), "mime_type": parts.get("mime_type"), "data": parts.get("data")}]
+            })
+
+        # Add new user input or images to the prompt
+        if user_input:
+            prompts.append({"role": "user", "parts": [{"text": user_input}]})
+        if uploaded_files:
+            for uploaded_file in uploaded_files:
+                image_base64 = get_image_base64(Image.open(uploaded_file))
+                prompts.append({"role": "user", "parts": [{"mime_type": uploaded_file.type, "data": image_base64}]})
+
         # Use Gemini Pro Vision model for image-based interaction
         vision_model = genai.GenerativeModel(
             model_name='gemini-pro-vision',
             generation_config=generation_config,
             safety_settings=safety_settings
         )
-        response = vision_model.generate_content(image_prompts)
-        response_text = response.text if hasattr(response, "text") else "No response text found."
-        
-        # Append images and response to chat history
-        for prompt in image_prompts:
-            st.session_state['chat_history'].append(prompt)
-        st.session_state['chat_history'].append({"role": "model", "parts": [{"text": response_text}]})
-        
-    # If no images are uploaded, use Gemini Pro model for text-based interaction
-    elif user_input:
-        text_model = genai.GenerativeModel(
-            model_name='gemini-pro',
-            generation_config=generation_config,
-            safety_settings=safety_settings
-        )
-        response = text_model.generate_content(st.session_state['chat_history'])
+        response = vision_model.generate_content(prompts)
         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}]})