Update app.py

app.py

@@ -4,8 +4,7 @@ import os
 from dotenv import load_dotenv
 import http.client
 import json
-import asyncio
-from typing import AsyncGenerator
+from typing import Iterator
 
 load_dotenv()
 
@@ -24,7 +23,7 @@ model = genai.GenerativeModel('gemini-2.0-flash-exp',
                             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")
 
-def perform_web_search(query):
+def perform_web_search(query: str) -> dict:
     conn = http.client.HTTPSConnection("google.serper.dev")
     payload = json.dumps({"q": query})
     headers = {
@@ -42,7 +41,7 @@ def perform_web_search(query):
     finally:
         conn.close()
 
-def format_search_results(data):
+def format_search_results(data: dict) -> str:
     if not data:
         return "Aucun résultat trouvé"
     
@@ -69,93 +68,106 @@ def format_search_results(data):
     
     return result
 
-def role_to_streamlit(role):
-    if role == "model":
-        return "assistant"
-    else:
-        return role
-
-# Initialize session state
-if "chat" not in st.session_state:
-    st.session_state.chat = model.start_chat(history=[])
-if "web_search" not in st.session_state:
-    st.session_state.web_search = False
-
-# Display Form Title
-st.title("Mariam AI!")
-
-# Settings section
-with st.sidebar:
-    st.title("Paramètres")
-    st.session_state.web_search = st.toggle("Activer la recherche web", value=st.session_state.web_search)
-
-# File upload section
-uploaded_file = st.file_uploader("Télécharger un fichier (image/document)", type=['jpg', 'mp4', 'mp3', 'jpeg', 'png', 'pdf', 'txt'])
+def stream_response(prompt: str, uploaded_file=None) -> Iterator[str]:
+    """Stream the response from Gemini"""
+    try:
+        if uploaded_file:
+            response = model.generate_content([uploaded_file, "\n\n", prompt], stream=True)
+        else:
+            response = model.generate_content(prompt, stream=True)
+        
+        for chunk in response:
+            if chunk.text:
+                yield chunk.text
+                
+    except Exception as e:
+        yield f"Erreur lors de la génération de la réponse : {str(e)}"
 
-# Display chat messages
-for message in st.session_state.chat.history:
-    with st.chat_message(role_to_streamlit(message.role)):
-        st.markdown(message.parts[0].text)
+def role_to_streamlit(role: str) -> str:
+    return "assistant" if role == "model" else role
 
-def process_uploaded_file(file):
+def process_uploaded_file(file) -> object:
     if file is not None:
-        with open(os.path.join("temp", file.name), "wb") as f:
+        file_path = os.path.join("temp", file.name)
+        with open(file_path, "wb") as f:
             f.write(file.getbuffer())
         try:
-            gemini_file = genai.upload_file(os.path.join("temp", file.name))
-            return gemini_file
+            return genai.upload_file(file_path)
         except Exception as e:
             st.error(f"Erreur lors du téléchargement du fichier : {e}")
             return None
+        finally:
+            # Clean up the temporary file
+            if os.path.exists(file_path):
+                os.remove(file_path)
+
+def main():
+    # Initialize session state
+    if "chat" not in st.session_state:
+        st.session_state.chat = model.start_chat(history=[])
+    if "web_search" not in st.session_state:
+        st.session_state.web_search = False
+
+    st.title("Mariam AI!")
+
+    # Settings sidebar
+    with st.sidebar:
+        st.title("Paramètres")
+        st.session_state.web_search = st.toggle("Activer la recherche web", 
+                                              value=st.session_state.web_search)
+
+    # File upload
+    uploaded_file = st.file_uploader("Télécharger un fichier (image/document)", 
+                                   type=['jpg', 'mp4', 'mp3', 'jpeg', 'png', 'pdf', 'txt'])
+
+    # Display chat history
+    for message in st.session_state.chat.history:
+        with st.chat_message(role_to_streamlit(message.role)):
+            st.markdown(message.parts[0].text)
+
+    # Chat input
+    if prompt := st.chat_input("Hey?"):
+        # Display user message
+        st.chat_message("user").markdown(prompt)
+
+        # Handle file upload
+        uploaded_gemini_file = None
+        if uploaded_file:
+            uploaded_gemini_file = process_uploaded_file(uploaded_file)
 
-async def stream_response(prompt: str, uploaded_gemini_file=None) -> AsyncGenerator[str, None]:
-    try:
-        if uploaded_gemini_file:
-            response = await st.session_state.chat.send_message_async([uploaded_gemini_file, "\n\n", prompt], stream=True)
-        else:
-            response = await st.session_state.chat.send_message_async(prompt, stream=True)
-        
-        async for chunk in response:
-            if chunk.text:
-                yield chunk.text
-    except Exception as e:
-        st.error(f"Erreur lors du streaming : {e}")
-        yield "Désolé, une erreur s'est produite lors de la génération de la réponse."
-
-# Chat input and processing
-if prompt := st.chat_input("Hey?"):
-    uploaded_gemini_file = None
-    if uploaded_file:
-        uploaded_gemini_file = process_uploaded_file(uploaded_file)
-    
-    # Display user message
-    st.chat_message("user").markdown(prompt)
-    
-    try:
-        # Perform web search if enabled
-        web_results = None
-        if st.session_state.web_search:
-            with st.spinner("Recherche web en cours..."):
-                web_results = perform_web_search(prompt)
-                if web_results:
-                    formatted_results = format_search_results(web_results)
-                    prompt = f"""Question: {prompt}\n\nRésultats de recherche web:\n{formatted_results}\n\nPourrais-tu analyser ces informations et me donner une réponse complète?"""
-        
-        # Create a placeholder for the streaming response
-        with st.chat_message("assistant"):
-            response_placeholder = st.empty()
-            full_response = ""
-            
-            # Stream the response
-            for response_chunk in asyncio.run(stream_response(prompt, uploaded_gemini_file)):
-                full_response += response_chunk
-                response_placeholder.markdown(full_response + "▌")
+        try:
+            # Perform web search if enabled
+            if st.session_state.web_search:
+                with st.spinner("Recherche web en cours..."):
+                    web_results = perform_web_search(prompt)
+                    if web_results:
+                        formatted_results = format_search_results(web_results)
+                        prompt = f"""Question: {prompt}\n\nRésultats de recherche web:\n{formatted_results}\n\nPourrais-tu analyser ces informations et me donner une réponse complète?"""
+
+            # Display assistant message with streaming
+            with st.chat_message("assistant"):
+                message_placeholder = st.empty()
+                full_response = ""
+                
+                # Stream the response
+                for chunk in stream_response(prompt, uploaded_gemini_file):
+                    full_response += chunk
+                    # Update the message placeholder with the accumulated response
+                    message_placeholder.markdown(full_response + "▌")
+                
+                # Remove the cursor and update with the final response
+                message_placeholder.markdown(full_response)
             
-            # Update the placeholder with the complete response
-            response_placeholder.markdown(full_response)
-    
-    except Exception as e:
-        st.error(f"Erreur lors de l'envoi du message : {e}")
+            # Add the interaction to chat history
+            st.session_state.chat.history.extend([
+                {"role": "user", "parts": [prompt]},
+                {"role": "model", "parts": [full_response]}
+            ])
+
+        except Exception as e:
+            st.error(f"Erreur lors de l'envoi du message : {e}")
 
-# Create temp directory
-os.makedirs("temp", exist_ok=True)
+if __name__ == "__main__":
+    # Create temp directory
+    os.makedirs("temp", exist_ok=True)
+    main()