Update app.py

app.py

@@ -23,8 +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 'last_model_used' not in st.session_state:
-    st.session_state['last_model_used'] = 'text'
+if 'use_vision_model' not in st.session_state:
+    st.session_state['use_vision_model'] = False
 
 # UI layout
 st.title("Gemini Chatbot")
@@ -41,29 +41,39 @@ def get_image_base64(image):
 def clear_conversation():
     st.session_state['chat_history'] = []
     st.session_state['file_uploader_key'] = str(uuid.uuid4())
-    st.session_state['last_model_used'] = 'text'
+    st.session_state['use_vision_model'] = False
 
 # Function to send message and clear input
 def send_message():
     user_input = st.session_state.user_input
     uploaded_files = st.session_state.uploaded_files
 
-    # Determine which model to use based on input type
-    model_name = 'gemini-pro-vision' if uploaded_files else 'gemini-pro'
-    st.session_state['last_model_used'] = 'vision' if uploaded_files else 'text'
+    # If images are uploaded, switch to using the vision model
+    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 prompts for conversation
-    prompts = []
+    # If there is text input, add it to the chat history string
     if user_input:
-        prompts.append({"role": "user", "parts": [{"text": user_input}]})
-    
-    # Append images to prompts if uploaded
+        chat_history_str += f"\n{user_input}"
+        st.session_state['chat_history'].append({"role": "user", "parts": [{"text": user_input}]})
+
+    # Add uploaded images to chat history string and session state
     if uploaded_files:
         for uploaded_file in uploaded_files:
-            prompts.append({"role": "user", "parts": [{"mime_type": uploaded_file.type, "data": get_image_base64(Image.open(uploaded_file))}]})
-    
-    # Create a new list combining chat history with current prompts
-    combined_prompts = st.session_state['chat_history'] + prompts
+            base64_image = get_image_base64(Image.open(uploaded_file))
+            chat_history_str += f"\n[Image]"
+            st.session_state['chat_history'].append({
+                "role": "user",
+                "parts": [{"mime_type": uploaded_file.type, "data": base64_image}]
+            })
+
+    # Determine which model to use based on whether an image has been uploaded
+    model_name = 'gemini-pro-vision' if st.session_state['use_vision_model'] else 'gemini-pro'
 
     # Use the appropriate model for interaction
     model = genai.GenerativeModel(
@@ -71,19 +81,34 @@ def send_message():
         generation_config=generation_config,
         safety_settings=safety_settings
     )
-    response = model.generate_content(combined_prompts)
+
+    # Generate content from the chat history or the latest prompt
+    response = model.generate_content([{"role": "user", "parts": [{"text": chat_history_str}]}])
     response_text = response.text if hasattr(response, "text") else "No response text found."
-    
-    # Update chat history and display the model response
-    for prompt in prompts:
-        st.session_state['chat_history'].append(prompt)
-    st.session_state['chat_history'].append({"role": "model", "parts": [{"text": response_text}]})
+
+    # Display the model response
+    if response_text:
+        st.session_state['chat_history'].append({"role": "model", "parts": [{"text": response_text}]})
 
     # Clear the user input and reset the file uploader widget
     st.session_state.user_input = ''
     st.session_state.uploaded_files = []
     st.session_state.file_uploader_key = str(uuid.uuid4())
 
+    # Display chat history
+    display_chat_history()
+
+# Function to display the chat history
+def display_chat_history():
+    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:
+            # Display the image
+            st.image(Image.open(io.BytesIO(base64.b64decode(parts['data']))), caption='Uploaded Image')
+
 # Multiline text input for the user to send messages
 user_input = st.text_area("Enter your message here:", key="user_input")
 
@@ -101,15 +126,5 @@ send_button = st.button("Send", on_click=send_message)
 # Button to clear the conversation
 clear_button = st.button("Clear Conversation", on_click=clear_conversation)
 
-# Display the chat history
-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:
-        # Display the image
-        st.image(Image.open(io.BytesIO(base64.b64decode(parts['data']))), caption='Uploaded Image')
-
 # Ensure the file_uploader widget state is tied to the randomly generated key
 st.session_state.uploaded_files = uploaded_files