Update app.py

app.py

@@ -2,12 +2,13 @@ import streamlit as st
 import google.generativeai as genai
 import os
 from dotenv import load_dotenv
+from PIL import Image
+import mimetypes
 
 load_dotenv()
 # Configure the API key
 genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
 
-
 safety_settings = [
     {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
     {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
@@ -15,7 +16,6 @@ safety_settings = [
     {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
 ]
 
-
 genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
 
 model = genai.GenerativeModel('gemini-1.5-flash',safety_settings=safety_settings,
@@ -26,14 +26,12 @@ model = genai.GenerativeModel('gemini-1.5-flash',safety_settings=safety_settings
 # Function to get response from the model
 # Gemini uses 'model' for assistant; Streamlit uses 'assistant'
 
-
 def role_to_streamlit(role):
     if role == "model":
         return "assistant"
     else:
         return role
 
-
 # Add a Gemini Chat history object to Streamlit session state
 if "chat" not in st.session_state:
     st.session_state.chat = model.start_chat(history=[])
@@ -44,15 +42,62 @@ st.title("Mariam AI!")
 # Display chat messages from history above current input box
 for message in st.session_state.chat.history:
     with st.chat_message(role_to_streamlit(message.role)):
-        st.markdown(message.parts[0].text)
+        # Check if the message part is text or a file part
+        for part in message.parts:
+            if part.HasField("text"):
+                st.markdown(part.text)
+            elif part.HasField("file_data"):
+                # Handle file display (e.g., image)
+                try:
+                    # Infer MIME type if not provided
+                    if not part.file_data.mime_type:
+                        mime_type = mimetypes.guess_type(part.file_data.file_name)[0]
+                    else:
+                        mime_type = part.file_data.mime_type
+                    if mime_type and mime_type.startswith("image/"):
+                        image_data = part.file_data.data  # Access the image data directly
+                        image = Image.open(io.BytesIO(image_data)) # Open the image using PIL
+                        st.image(image)
+                    else:
+                        st.write(f"File: {part.file_data.file_name} (MIME type: {part.file_data.mime_type})")
+                except Exception as e:
+                    st.error(f"Error displaying file: {e}")
 
 # Accept user's next message, add to context, resubmit context to Gemini
 if prompt := st.chat_input("Hey?"):
     # Display user's last message
     st.chat_message("user").markdown(prompt)
 
-    # Send user entry to Gemini and read the response
-    response = st.session_state.chat.send_message(prompt)
+    # Handle file uploads
+    uploaded_file = st.file_uploader("Choose a file", type=["jpg", "jpeg", "png", "pdf"]) # Add more type if needed
+
+    if uploaded_file is not None:
+        # To read file as bytes:
+        bytes_data = uploaded_file.getvalue()
+
+        # Display the uploaded file (if it's an image)
+        if uploaded_file.type.startswith("image/"):
+            image = Image.open(uploaded_file)
+            st.image(image, caption=f"Uploaded Image: {uploaded_file.name}")
+
+        # Construct the message parts, including the uploaded file
+        parts = [
+            prompt,
+            {
+                "file_data": {
+                    "mime_type": uploaded_file.type,
+                    "file_name": uploaded_file.name,
+                    "data": bytes_data
+                }
+            }
+        ]
+
+        # Send the message with file to Gemini
+        response = st.session_state.chat.send_message(parts)
+    
+    else:
+      # Send user entry to Gemini and read the response
+      response = st.session_state.chat.send_message(prompt)
 
     # Display last
     with st.chat_message("assistant"):