Update app.py

app.py

@@ -4,10 +4,9 @@ import io
 import base64
 import uuid
 
-# Assuming google.generativeai is the correct import based on your description
+# Assuming google.generativeai is correctly imported as genai and the API key is set
 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
 genai.configure(api_key=api_key)
 
@@ -18,7 +17,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:
@@ -26,33 +24,25 @@ if 'file_uploader_key' not in st.session_state:
 if 'use_vision_model' not in st.session_state:
     st.session_state['use_vision_model'] = False
 
-# UI layout
 st.title("Gemini Chatbot")
 
-# Function to convert image to base64
 def get_image_base64(image):
-    image = image.convert("RGB")  # Convert to RGB to remove alpha channel if present
+    image = image.convert("RGB")
     buffered = io.BytesIO()
     image.save(buffered, format="JPEG")
     img_str = base64.b64encode(buffered.getvalue()).decode()
     return img_str
 
-# Function to clear conversation
 def clear_conversation():
     st.session_state['chat_history'] = []
     st.session_state['file_uploader_key'] = str(uuid.uuid4())
     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
-
-    # If images are uploaded, switch to using the vision model
     if uploaded_files:
         st.session_state['use_vision_model'] = True
-
-    # Prepare the prompt for the vision model
     prompts = []
     for entry in st.session_state['chat_history']:
         for part in entry['parts']:
@@ -60,13 +50,9 @@ def send_message():
                 prompts.append(part['text'])
             elif 'data' in part:
                 prompts.append("[Image]")
-
-    # If there is text input, add it to the list of prompts
     if user_input:
         prompts.append(user_input)
         st.session_state['chat_history'].append({"role": "user", "parts": [{"text": user_input}]})
-
-    # 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))
@@ -75,23 +61,14 @@ def send_message():
                 "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(
         model_name=model_name,
         generation_config=generation_config,
         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
     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']
@@ -99,23 +76,15 @@ def send_message():
         ]
     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
     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"]
@@ -123,13 +92,19 @@ def display_chat_history():
         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
+def get_chat_history_str():
+    chat_history_str = "\n".join(
+        f"{entry['role'].title()}: {part['text']}" if 'text' in part 
+        else f"{entry['role'].title()}: (Image)"
+        for entry in st.session_state['chat_history'] 
+        for part in entry['parts']
+    )
+    return chat_history_str
+
 user_input = st.text_area("Enter your message here:", key="user_input")
 
-# File uploader for images
 uploaded_files = st.file_uploader(
     "Upload images:",
     type=["png", "jpg", "jpeg"],
@@ -137,11 +112,25 @@ uploaded_files = st.file_uploader(
     key=st.session_state.file_uploader_key
 )
 
-# Button to send the message
 send_button = st.button("Send", on_click=send_message)
 
-# Button to clear the conversation
 clear_button = st.button("Clear Conversation", on_click=clear_conversation)
 
+# Function to download the chat history
+def download_chat_history():
+    chat_history_str = get_chat_history_str()
+    return chat_history_str
+
+# Add a button to download the chat history as a text file
+download_button = st.download_button(
+    label="Download Chat",
+    data=download_chat_history(),
+    file_name="chat_history.txt",
+    mime="text/plain"
+)
+
+# Display the chat history
+display_chat_history()
+
 # 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