Update app.py

app.py

@@ -19,8 +19,6 @@ generation_config = genai.GenerationConfig(
 safety_settings = []
 
 # Initialize session state for chat history and file uploader key
-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:
@@ -39,7 +37,6 @@ 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())
 
@@ -48,27 +45,25 @@ def send_message():
     user_input = st.session_state.user_input
     uploaded_files = st.session_state.uploaded_files
 
-    # If images are uploaded or we have already used the vision model, 
-    # bundle the entire conversation history into a single-turn interaction.
+    # Check if an image has been uploaded or if we are continuing with the vision model
     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}]})
+        image_prompt = None
+
+        # Create the prompt for the vision model
         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}]})
+            image = Image.open(uploaded_files[0])  # Only take the first image for simplicity
+            image_base64 = get_image_base64(image)
+            image_prompt = {
+                "role": "user",
+                "parts": [{"mime_type": uploaded_files[0].type, "data": image_base64}]
+            }
+        elif user_input:
+            # Text input after using vision model
+            image_prompt = {
+                "role": "user",
+                "parts": [{"text": user_input}]
+            }
 
         # Use Gemini Pro Vision model for image-based interaction
         vision_model = genai.GenerativeModel(
@@ -76,14 +71,31 @@ def send_message():
             generation_config=generation_config,
             safety_settings=safety_settings
         )
-        response = vision_model.generate_content(prompts)
+
+        response = vision_model.generate_content([image_prompt])
+        response_text = response.text if hasattr(response, "text") else "No response text found."
+        st.write("AI: " + response_text)
+
+    # If no images are uploaded and we haven't used the vision model yet, use Gemini Pro model
+    elif user_input and not st.session_state['using_vision_model']:
+        text_prompt = {
+            "role": "user",
+            "parts": [{"text": user_input}]
+        }
+
+        text_model = genai.GenerativeModel(
+            model_name='gemini-pro',
+            generation_config=generation_config,
+            safety_settings=safety_settings
+        )
+
+        response = text_model.generate_content([text_prompt])
         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}]})
+        st.write("AI: " + response_text)
 
-    # Clear the user input and generate a new key for the file uploader widget to reset it
+    # Clear the user input and uploaded files
     st.session_state.user_input = ''
     st.session_state.uploaded_files = []
-    st.session_state.file_uploader_key = str(uuid.uuid4())
 
 # Multiline text input for the user to send messages
 user_input = st.text_area("Enter your message here:", key="user_input")
@@ -92,7 +104,7 @@ user_input = st.text_area("Enter your message here:", key="user_input")
 uploaded_files = st.file_uploader(
     "Upload images:",
     type=["png", "jpg", "jpeg"],
-    accept_multiple_files=True,
+    accept_multiple_files=False,  # For simplicity, we're only accepting one image at a time
     key=st.session_state.file_uploader_key
 )
 
@@ -100,16 +112,4 @@ uploaded_files = st.file_uploader(
 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:
-        st.markdown(f"{role.title()}: (Image)")
-
-# Ensure the file_uploader widget state is tied to the randomly generated key
-st.session_state.uploaded_files = uploaded_files
+clear_button = st.button("Clear Conversation", on_click=clear_conversation)