Update app.py

app.py

@@ -8,7 +8,7 @@ import uuid
 import google.generativeai as genai
 
 # Configure the API key (should be set as an environment variable or secure storage in production)
-api_key = "AIzaSyC70u1sN87IkoxOoIj4XCAPw97ae2LZwNM"  # Replace with your actual API key
+api_key = "your_actual_api_key"  # Replace with your actual API key
 genai.configure(api_key=api_key)
 
 generation_config = genai.GenerationConfig(
@@ -19,10 +19,12 @@ 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:
-    st.session_state['using_vision_model'] = False
+if 'last_model_used' not in st.session_state:
+    st.session_state['last_model_used'] = 'text'
 
 # UI layout
 st.title("Gemini Chatbot")
@@ -37,65 +39,50 @@ def get_image_base64(image):
 
 # Function to clear conversation
 def clear_conversation():
-    st.session_state['using_vision_model'] = False
+    st.session_state['chat_history'] = []
     st.session_state['file_uploader_key'] = str(uuid.uuid4())
+    st.session_state['last_model_used'] = 'text'
 
 # Function to send message and clear input
 def send_message():
     user_input = st.session_state.user_input
     uploaded_files = st.session_state.uploaded_files
 
-    # 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
-        image_prompt = None
-
-        # Create the prompt for the vision model
-        if uploaded_files:
-            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(
-            model_name='gemini-pro-vision',
-            generation_config=generation_config,
-            safety_settings=safety_settings
-        )
-
-        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.write("AI: " + response_text)
-
-    # Clear the user input and 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'
+
+    # Prepare prompts for conversation
+    prompts = []
+    if user_input:
+        prompts.append({"role": "user", "parts": [{"text": user_input}]})
+    
+    # Append images to prompts if uploaded
+    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
+
+    # Use the appropriate model for interaction
+    model = genai.GenerativeModel(
+        model_name=model_name,
+        generation_config=generation_config,
+        safety_settings=safety_settings
+    )
+    response = model.generate_content(combined_prompts)
+    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}]})
+
+    # 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())
 
 # Multiline text input for the user to send messages
 user_input = st.text_area("Enter your message here:", key="user_input")
@@ -104,7 +91,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=False,  # For simplicity, we're only accepting one image at a time
+    accept_multiple_files=True,
     key=st.session_state.file_uploader_key
 )
 
@@ -112,4 +99,17 @@ 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)
+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