Update app.py

app.py

@@ -3,6 +3,9 @@ import google.generativeai as genai
 import os
 from dotenv import load_dotenv
 from PIL import Image
+import tempfile
+import time
+import ssl
 
 # Charger les variables d'environnement
 load_dotenv()
@@ -18,56 +21,121 @@ safety_settings = [
     {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
 ]
 
+def upload_and_process_file(file_path):
+    max_retries = 3
+    retry_delay = 2
+    
+    for attempt in range(max_retries):
+        try:
+            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")
+                
+            uploaded_file = genai.upload_file(path=file_path)
+            
+            timeout = 300
+            start_time = time.time()
+            
+            while uploaded_file.state.name == "PROCESSING":
+                if time.time() - start_time > timeout:
+                    raise TimeoutError("Timeout pendant le traitement du fichier")
+                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}")
+                
+            return uploaded_file
+
+        except Exception as e:
+            if attempt < max_retries - 1:
+                time.sleep(retry_delay * (attempt + 1))
+            else:
+                raise
+
+def allowed_file(filename):
+    ALLOWED_EXTENSIONS = {'txt', 'pdf', 'png', 'jpg', 'jpeg', 'gif'}
+    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
+
 # Initialiser le modèle
 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")
-
-# Fonction pour convertir le rôle
-def role_to_streamlit(role):
-    return "assistant" if role == "model" else role
+                            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")
 
 # Configuration de la page Streamlit
 st.set_page_config(page_title="Mariam - Assistant IA", page_icon="🤖")
 st.title("Mariam AI - Chat Intelligent")
 
-# Initialiser l'historique de chat si pas déjà présent
+# Initialiser l'historique de chat
 if "chat" not in st.session_state:
     st.session_state.chat = model.start_chat(history=[])
 
-# Gestion de l'upload d'images
-uploaded_file = st.file_uploader("Télécharger une image (optionnel)", 
-                                  type=["jpg", "jpeg", "png", "gif", "bmp"])
+# Créer trois colonnes pour les boutons d'upload
+col1, col2, col3 = st.columns([1, 1, 8])
+
+# Bouton pour les images
+with col1:
+    uploaded_images = st.file_uploader("", type=["jpg", "jpeg", "png", "gif"], 
+                                     accept_multiple_files=True, key="images")
+    st.markdown("📸")
+
+# Bouton pour les autres fichiers
+with col2:
+    uploaded_files = st.file_uploader("", type=["txt", "pdf"], 
+                                    accept_multiple_files=True, key="files")
+    st.markdown("📁")
 
 # Afficher l'historique des messages
 for message in st.session_state.chat.history:
     with st.chat_message(role_to_streamlit(message.role)):
         st.markdown(message.parts[0].text)
-        # Afficher l'image si présente dans l'historique
-        if len(message.parts) > 1 and hasattr(message.parts[1], 'image'):
-            st.image(message.parts[1].image)
+        if len(message.parts) > 1:
+            for part in message.parts[1:]:
+                if hasattr(part, 'image'):
+                    st.image(part.image)
 
 # Zone de saisie du chat
 if prompt := st.chat_input("Que puis-je faire pour vous ?"):
-    # Préparer le message
-    if uploaded_file is not None:
-        # Convertir le fichier téléchargé en image
-        image = Image.open(uploaded_file)
-        
-        # Envoyer le message avec l'image
-        response = st.session_state.chat.send_message([
-            prompt, 
-            image
-        ])
-        
-        # Afficher l'image téléchargée
-        st.chat_message("user").markdown(prompt)
-        st.chat_message("user").image(image)
-    else:
-        # Envoyer le message texte seul
-        response = st.session_state.chat.send_message(prompt)
+    content = [prompt]
+    temp_files = []
+    
+    try:
+        # Traitement des images
+        if uploaded_images:
+            for img_file in uploaded_images:
+                if allowed_file(img_file.name):
+                    image = Image.open(img_file)
+                    content.append(image)
+                    st.chat_message("user").image(image)
+
+        # Traitement des autres fichiers
+        if uploaded_files:
+            for file in uploaded_files:
+                if allowed_file(file.name):
+                    with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.name)[1]) as temp_file:
+                        temp_file.write(file.getvalue())
+                        temp_files.append(temp_file.name)
+                        uploaded_file = upload_and_process_file(temp_file.name)
+                        content.append(uploaded_file)
+
+        # Afficher le message utilisateur
         st.chat_message("user").markdown(prompt)
+        
+        # Envoyer le message et afficher la réponse
+        response = st.session_state.chat.send_message(content)
+        with st.chat_message("assistant"):
+            st.markdown(response.text)
 
-    # Afficher la réponse de l'assistant
-    with st.chat_message("assistant"):
-        st.markdown(response.text)
+    except Exception as e:
+        st.error(f"Une erreur est survenue : {str(e)}")
+    
+    finally:
+        # Nettoyage des fichiers temporaires
+        for temp_file in temp_files:
+            try:
+                os.unlink(temp_file)
+            except Exception as e:
+                print(f"Erreur lors de la suppression du fichier temporaire {temp_file}: {e}")