Update app.py

app.py

@@ -52,21 +52,25 @@ def send_message():
     if uploaded_files:
         st.session_state['use_vision_model'] = True
 
-    # Prepare the prompt with chat history for the vision model
-    chat_history_str = "\n".join(
-        part['text'] for entry in st.session_state['chat_history'] for part in entry['parts'] if 'text' in part
-    )
+    # Prepare the prompt for the vision model
+    prompts = []
+    for entry in st.session_state['chat_history']:
+        for part in entry['parts']:
+            if 'text' in part:
+                prompts.append(part['text'])
+            elif 'data' in part:
+                prompts.append("[Image]")
 
-    # If there is text input, add it to the chat history string
+    # If there is text input, add it to the list of prompts
     if user_input:
-        chat_history_str += f"\n{user_input}"
+        prompts.append(user_input)
         st.session_state['chat_history'].append({"role": "user", "parts": [{"text": user_input}]})
 
-    # Add uploaded images to chat history string and session state
+    # Add uploaded images to the list of prompts and session state
     if uploaded_files:
         for uploaded_file in uploaded_files:
             base64_image = get_image_base64(Image.open(uploaded_file))
-            chat_history_str += f"\n[Image]"
+            prompts.append("[Image]")
             st.session_state['chat_history'].append({
                 "role": "user",
                 "parts": [{"mime_type": uploaded_file.type, "data": base64_image}]
@@ -82,8 +86,21 @@ def send_message():
         safety_settings=safety_settings
     )
 
+    # Create a single prompt string from the list of prompts
+    chat_history_str = "\n".join(prompts)
+
     # Generate content from the chat history or the latest prompt
-    response = model.generate_content([{"role": "user", "parts": [{"text": chat_history_str}]}])
+    if st.session_state['use_vision_model']:
+        # When using vision model, include images as base64 strings
+        prompt_parts = [{"text": chat_history_str}] + [
+            {"data": part['data'], "mime_type": "image/jpeg"}
+            for entry in st.session_state['chat_history'] for part in entry['parts']
+            if 'data' in part
+        ]
+    else:
+        prompt_parts = [{"text": chat_history_str}]
+
+    response = model.generate_content([{"role": "user", "parts": prompt_parts}])
     response_text = response.text if hasattr(response, "text") else "No response text found."
 
     # Display the model response
@@ -127,4 +144,4 @@ send_button = st.button("Send", on_click=send_message)
 clear_button = st.button("Clear Conversation", on_click=clear_conversation)
 
 # Ensure the file_uploader widget state is tied to the randomly generated key
-st.session_state.uploaded_files = uploaded_files
+st.session_state.uploaded_files = uploaded_files