Update app.py

app.py

@@ -1,14 +1,11 @@
 import streamlit as st
 import google.generativeai as genai
 import os
-import tempfile
-import PIL.Image
-import time
-import ssl
 from dotenv import load_dotenv
+from mimetypes import guess_type as guessmime
+from io import BytesIO
 
 load_dotenv()
-
 # Configure the API key
 genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
 
@@ -19,16 +16,36 @@ safety_settings = [
     {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
 ]
 
-model = genai.GenerativeModel('gemini-1.5-flash-002', safety_settings=safety_settings,
-                              system_instruction="Tu es un assistant intelligent. ton but est d'assister au mieux que tu peux. tu as été créé par Aenir et tu t'appelles Mariam")
+genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
+
+model = genai.GenerativeModel('gemini-1.5-flash',safety_settings=safety_settings,
+            system_instruction="Tu es un assistant intelligent. ton but est d'assister au mieux que tu peux. tu as été créé par Aenir et tu t'appelles Mariam")
+
+
 
 # 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
+    
+def create_media_part(data, mimetype,filename=None):
+        """Creates a media part for the GenerativeModel.
+
+        Args:
+            data: The image data as bytes.
+            mimetype: The mimetype of the image.
+            filename: optional filename for the image part
+
+        Returns:
+            A Part object representing the image.
+        """
+        # The API expects raw bytes so avoid having it base64 encode the input
+        blob = genai.types.Blob(mimetype, genai.types.bytes_to_data(data, mime_type=mimetype))
+        return genai.types.Part(filename=filename, inline_data=blob)
 
 # Add a Gemini Chat history object to Streamlit session state
 if "chat" not in st.session_state:
@@ -37,113 +54,41 @@ if "chat" not in st.session_state:
 # Display Form Title
 st.title("Mariam AI!")
 
-# File uploader
-uploaded_files = st.file_uploader("Choose a file", accept_multiple_files=True)
-
 # 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)
 
-def upload_and_process_file(file_path):
-    """Upload et traite un fichier avec l'API Gemini avec gestion des erreurs améliorée"""
-    max_retries = 3
-    retry_delay = 2  # secondes
-
-    for attempt in range(max_retries):
-        try:
-            print(f"Tentative d'upload {attempt + 1}/{max_retries} pour {file_path}")
-
-            # Vérification du fichier
-            if not os.path.exists(file_path):
-                raise FileNotFoundError(f"Le fichier {file_path} n'existe pas")
-
-            file_size = os.path.getsize(file_path)
-            if file_size == 0:
-                raise ValueError(f"Le fichier {file_path} est vide")
-
-            # Upload du fichier
-            uploaded_file = genai.upload_file(path=file_path)
-            print(f"Upload réussi: {uploaded_file.uri}")
-
-            # Attente du traitement
-            timeout = 300  # 5 minutes
-            start_time = time.time()
-
-            while uploaded_file.state.name == "PROCESSING":
-                if time.time() - start_time > timeout:
-                    raise TimeoutError("Timeout pendant le traitement du fichier")
-
-                print(
-                    f"En attente du traitement... Temps écoulé: {int(time.time() - start_time)}s")
-                time.sleep(10)
-                uploaded_file = genai.get_file(uploaded_file.name)
-
-            if uploaded_file.state.name == "FAILED":
-                raise ValueError(
-                    f"Échec du traitement: {uploaded_file.state.name}")
-
-            print(f"Traitement terminé avec succès: {uploaded_file.uri}")
-            return uploaded_file
-
-        except ssl.SSLError as e:
-            print(
-                f"Erreur SSL lors de l'upload (tentative {attempt + 1}): {e}")
-            if attempt < max_retries - 1:
-                time.sleep(retry_delay * (attempt + 1))
-            else:
-                raise
-
-        except Exception as e:
-            print(
-                f"Erreur lors de l'upload (tentative {attempt + 1}): {e}")
-            if attempt < max_retries - 1:
-                time.sleep(retry_delay * (attempt + 1))
-            else:
-                raise
-
 # 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)
 
-    content = [prompt]
-    temp_files = []
-
-    try:
-        # Process uploaded files
-        for uploaded_file in uploaded_files:
-            with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]) as temp_file:
-                temp_file.write(uploaded_file.getvalue())
-                temp_files.append(temp_file.name)
-                if uploaded_file.name.lower().endswith(('.png', '.jpg', '.jpeg', '.gif')):
-                    content.append(PIL.Image.open(temp_file.name))
-                else:
-                    processed_file = upload_and_process_file(temp_file.name)
-                    content.append(processed_file)
-
+    # Handle file uploads
+    uploaded_file = st.file_uploader("Choose a file", type=["png", "jpg", "jpeg", "mp3", "wav", "mp4", "avi"])
+
+    if uploaded_file is not None:
+        # Display the uploaded file
+        if uploaded_file.type.startswith('image'):
+            st.image(uploaded_file, caption="Uploaded Image.", use_column_width=True)
+        elif uploaded_file.type.startswith('audio'):
+            st.audio(uploaded_file, format=uploaded_file.type)
+        elif uploaded_file.type.startswith('video'):
+            st.video(uploaded_file, format=uploaded_file.type)
+    
+        file_bytes = uploaded_file.getvalue()
+        mime_type = guessmime(uploaded_file.name)[0]
+        media_part = create_media_part(file_bytes, mime_type,filename=uploaded_file.name)
         
+        # Combine prompt and uploaded file for Gemini
+        parts = [media_part, "\n\n", prompt]
+        response = st.session_state.chat.send_message(parts)
+
+    else:
+
         # Send user entry to Gemini and read the response
-        response = model.generate_content(content, stream=True)
-        response.resolve()
-
-        # Display the response
-        with st.chat_message("assistant"):
-            st.markdown(response.text)
-
-        # Update the chat history
-        st.session_state.chat.history.extend([
-            genai.types.Content(parts=[genai.types.Part(text=prompt)], role="user"),
-            genai.types.Content(parts=[genai.types.Part(text=response.text)], role="model")
-        ])
-
-    except Exception as e:
-        st.error(f"An error occurred: {e}")
-
-    finally:
-        # Cleanup temporary files
-        for temp_file in temp_files:
-            try:
-                os.unlink(temp_file)
-            except Exception as e:
-                print(f"Error deleting temporary file {temp_file}: {e}")
+        response = st.session_state.chat.send_message(prompt)
+
+    # Display last
+    with st.chat_message("assistant"):
+        st.markdown(response.text)