Update app.py

app.py

@@ -55,23 +55,20 @@ selected_model = st.selectbox("Choose a Gemini 1.5 Model:", ["gemini-1.5-flash-l
 enable_tts = st.checkbox("Enable Text-to-Speech")
 
 # --- Helper Functions ---
-def get_image_base64(image):
-    image = image.convert("RGB")
-    buffered = io.BytesIO()
-    image.save(buffered, format="JPEG")
-    img_str = base64.b64encode(buffered.getvalue()).decode()
-    return img_str
+def get_file_base64(file_content, mime_type):
+    base64_data = base64.b64encode(file_content).decode()
+    return {"mime_type": mime_type, "data": base64_data}
 
 def clear_conversation():
     st.session_state['chat_history'] = []
     st.session_state['file_uploader_key'] = str(uuid.uuid4())
 
 def display_chat_history():
-    chat_container = st.empty()  # Create an empty container to hold the chat history
-    with chat_container.container():  # Use a container to update the chat history
+    chat_container = st.empty()
+    with chat_container.container():
         for entry in st.session_state['chat_history']:
             role = entry["role"]
-            parts = entry["parts"][0]
+            parts = entry["parts"][0] 
             if 'text' in parts:
                 st.markdown(f"**{role.title()}:** {parts['text']}")
             elif 'data' in parts:
@@ -85,6 +82,8 @@ def display_chat_history():
                     for page_num in range(len(pdf_reader.pages)):
                         page = pdf_reader.pages[page_num]
                         st.write(page.extract_text())
+                elif mime_type.startswith('audio'):
+                    st.audio(io.BytesIO(base64.b64decode(parts['data'])), format=mime_type)
                 elif mime_type.startswith('video'):
                     st.video(io.BytesIO(base64.b64decode(parts['data'])))
 
@@ -103,28 +102,33 @@ def send_message():
     if uploaded_files:
         for uploaded_file in uploaded_files:
             file_content = uploaded_file.read()
-            base64_data = base64.b64encode(file_content).decode()
-            prompt_parts.append({"mime_type": uploaded_file.type, "data": base64_data})
-            st.session_state['chat_history'].append({"role": "user", "parts": [{"mime_type": uploaded_file.type, "data": base64_data}]})
+            prompt_parts.append(get_file_base64(file_content, uploaded_file.type))
+            st.session_state['chat_history'].append(
+                {"role": "user", "parts": [get_file_base64(file_content, uploaded_file.type)]}
+            )
 
     # Generate response using the selected model
-    model = genai.GenerativeModel(
-        model_name=selected_model,
-        generation_config=generation_config,
-        safety_settings=safety_settings
-    )
-
-    response = model.generate_content([{"role": "user", "parts": prompt_parts}])
-    response_text = response.text if hasattr(response, "text") else "No response text found."
-
-    if response_text:
-        st.session_state['chat_history'].append({"role": "model", "parts": [{"text": response_text}]})
-        if enable_tts:
-            tts = gTTS(text=response_text, lang='en')
-            tts_file = BytesIO()
-            tts.write_to_fp(tts_file)
-            tts_file.seek(0)
-            st.audio(tts_file, format='audio/mp3')
+    try:
+        model = genai.GenerativeModel(
+            model_name=selected_model,
+            generation_config=generation_config,
+            safety_settings=safety_settings
+        )
+
+        response = model.generate_content([{"role": "user", "parts": prompt_parts}])
+        response_text = response.text if hasattr(response, "text") else "No response text found."
+
+        if response_text:
+            st.session_state['chat_history'].append({"role": "model", "parts": [{"text": response_text}]})
+            if enable_tts:
+                tts = gTTS(text=response_text, lang='en')
+                tts_file = BytesIO()
+                tts.write_to_fp(tts_file)
+                tts_file.seek(0)
+                st.audio(tts_file, format='audio/mp3')
+
+    except Exception as e:
+        st.error(f"An error occurred: {e}") 
 
     st.session_state.user_input = ''
     st.session_state.uploaded_files = []
@@ -151,8 +155,8 @@ with col2:
 
 # --- File Uploader ---
 uploaded_files = st.file_uploader(
-    "Upload Files (Images, Videos, PDFs):",
-    type=["png", "jpg", "jpeg", "mp4", "pdf"],
+    "Upload Files (Images, Videos, PDFs, MP3):",
+    type=["png", "jpg", "jpeg", "mp4", "pdf", "mp3"],
     accept_multiple_files=True,
     key=st.session_state.file_uploader_key
 )