Update app.py

app.py

@@ -340,22 +340,28 @@ async def update_audio_preview(audio_file):
 # Async function to generate lecture materials and audio
 async def on_generate(api_service, api_key, serpapi_key, title, topic, instructions, lecture_type, speaker_audio, num_slides):
     if not serpapi_key:
-        yield f"""
-        <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
-            <h2 style="color: #d9534f;">SerpApi key required</h2>
-            <p style="margin-top: 20px;">Please provide a valid SerpApi key and try again.</p>
-        </div>
-        """
+        yield {
+            "html": f"""
+            <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
+                <h2 style="color: #d9534f;">SerpApi key required</h2>
+                <p style="margin-top: 20px;">Please provide a valid SerpApi key and try again.</p>
+            </div>
+            """,
+            "files": []
+        }
         return
 
     if tts is None:
-        yield f"""
-        <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
-            <h2 style="color: #d9534f;">TTS model initialization failed</h2>
-            <p style="margin-top: 20px;">The TTS model failed to initialize at startup.</p>
-            <p>Please ensure the Coqui TTS model is properly installed and try restarting the application.</p>
-        </div>
-        """
+        yield {
+            "html": f"""
+            <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
+                <h2 style="color: #d9534f;">TTS model initialization failed</h2>
+                <p style="margin-top: 20px;">The TTS model failed to initialize at startup.</p>
+                <p>Please ensure the Coqui TTS model is properly installed and try restarting the application.</p>
+            </div>
+            """,
+            "files": []
+        }
         return
 
     model_client = get_model_client(api_service, api_key)
@@ -418,7 +424,10 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
     
     progress = 0
     label = "Research: in progress..."
-    yield html_with_progress(label, progress)
+    yield {
+        "html": html_with_progress(label, progress),
+        "files": []
+    }
     await asyncio.sleep(0.1)
     
     initial_message = f"""
@@ -458,7 +467,10 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
                 if source == "research_agent" and message.target == "slide_agent":
                     progress = 25
                     label = "Slides: generating..."
-                    yield html_with_progress(label, progress)
+                    yield {
+                        "html": html_with_progress(label, progress),
+                        "files": []
+                    }
                     await asyncio.sleep(0.1)
                 elif source == "slide_agent" and message.target == "script_agent":
                     if slides is None:
@@ -480,7 +492,10 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
                             continue
                     progress = 50
                     label = "Scripts: generating..."
-                    yield html_with_progress(label, progress)
+                    yield {
+                        "html": html_with_progress(label, progress),
+                        "files": []
+                    }
                     await asyncio.sleep(0.1)
                 elif source == "script_agent" and message.target == "feynman_agent":
                     if scripts is None:
@@ -491,14 +506,20 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
                             logger.info("Extracted scripts JSON from HandoffMessage context: %s", scripts)
                     progress = 75
                     label = "Review: in progress..."
-                    yield html_with_progress(label, progress)
+                    yield {
+                        "html": html_with_progress(label, progress),
+                        "files": []
+                    }
                     await asyncio.sleep(0.1)
             
             elif source == "research_agent" and isinstance(message, TextMessage) and "handoff_to_slide_agent" in message.content:
                 logger.info("Research Agent completed research")
                 progress = 25
                 label = "Slides: generating..."
-                yield html_with_progress(label, progress)
+                yield {
+                    "html": html_with_progress(label, progress),
+                    "files": []
+                }
                 await asyncio.sleep(0.1)
             
             elif source == "slide_agent" and isinstance(message, (TextMessage, StructuredMessage)):
@@ -528,7 +549,10 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
                             logger.error("Error saving slide content to %s: %s", content_file, str(e))
                     progress = 50
                     label = "Scripts: generating..."
-                    yield html_with_progress(label, progress)
+                    yield {
+                        "html": html_with_progress(label, progress),
+                        "files": []
+                    }
                     await asyncio.sleep(0.1)
                 else:
                     logger.warning("No JSON extracted from slide_agent message")
@@ -559,7 +583,10 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
                             logger.error("Error saving raw script to %s: %s", script_file, str(e))
                     progress = 75
                     label = "Scripts generated and saved. Reviewing..."
-                    yield html_with_progress(label, progress)
+                    yield {
+                        "html": html_with_progress(label, progress),
+                        "files": []
+                    }
                     await asyncio.sleep(0.1)
                 else:
                     logger.warning("No JSON extracted from script_agent message")
@@ -578,7 +605,10 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
                 logger.info("Feynman Agent completed lecture review: %s", message.content)
                 progress = 90
                 label = "Lecture materials ready. Generating audio..."
-                yield html_with_progress(label, progress)
+                yield {
+                    "html": html_with_progress(label, progress),
+                    "files": []
+                }
                 await asyncio.sleep(0.1)
         
         logger.info("Slides state: %s", "Generated" if slides else "None")
@@ -591,65 +621,83 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
             for msg in task_result.messages:
                 source = getattr(msg, 'source', getattr(msg, 'sender', None))
                 logger.debug("Message from %s, type: %s, content: %s", source, type(msg), msg.to_text() if hasattr(msg, 'to_text') else str(msg))
-            yield f"""
-            <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
-                <h2 style="color: #d9534f;">{error_message}</h2>
-                <p style="margin-top: 20px;">Please try again with a different model (e.g., Anthropic-claude-3-sonnet-20240229) or simplify the topic/instructions.</p>
-            </div>
-            """
+            yield {
+                "html": f"""
+                <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
+                    <h2 style="color: #d9534f;">{error_message}</h2>
+                    <p style="margin-top: 20px;">Please try again with a different model (e.g., Anthropic-claude-3-sonnet-20240229) or simplify the topic/instructions.</p>
+                </div>
+                """,
+                "files": []
+            }
             return
         
         if len(slides) != total_slides:
             logger.error("Expected %d slides, but received %d", total_slides, len(slides))
-            yield f"""
-            <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
-                <h2 style="color: #d9534f;">Incorrect number of slides</h2>
-                <p style="margin-top: 20px;">Expected {total_slides} slides, but generated {len(slides)}. Please try again.</p>
-            </div>
-            """
+            yield {
+                "html": f"""
+                <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
+                    <h2 style="color: #d9534f;">Incorrect number of slides</h2>
+                    <p style="margin-top: 20px;">Expected {total_slides} slides, but generated {len(slides)}. Please try again.</p>
+                </div>
+                """,
+                "files": []
+            }
             return
         
         if not isinstance(scripts, list) or not all(isinstance(s, str) for s in scripts):
             logger.error("Scripts are not a list of strings: %s", scripts)
-            yield f"""
-            <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
-                <h2 style="color: #d9534f;">Invalid script format</h2>
-                <p style="margin-top: 20px;">Scripts must be a list of strings. Please try again.</p>
-            </div>
-            """
+            yield {
+                "html": f"""
+                <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
+                    <h2 style="color: #d9534f;">Invalid script format</h2>
+                    <p style="margin-top: 20px;">Scripts must be a list of strings. Please try again.</p>
+                </div>
+                """,
+                "files": []
+            }
             return
         
         if len(scripts) != total_slides:
             logger.error("Mismatch between number of slides (%d) and scripts (%d)", len(slides), len(scripts))
-            yield f"""
-            <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
-                <h2 style="color: #d9534f;">Mismatch in slides and scripts</h2>
-                <p style="margin-top: 20px;">Generated {len(slides)} slides but {len(scripts)} scripts. Please try again.</p>
-            </div>
-            """
+            yield {
+                "html": f"""
+                <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
+                    <h2 style="color: #d9534f;">Mismatch in slides and scripts</h2>
+                    <p style="margin-top: 20px;">Generated {len(slides)} slides but {len(scripts)} scripts. Please try again.</p>
+                </div>
+                """,
+                "files": []
+            }
             return
         
         markdown_slides = generate_markdown_slides(slides, title)
         if not markdown_slides:
             logger.error("Failed to generate Markdown slides")
-            yield f"""
-            <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
-                <h2 style="color: #d9534f;">Failed to generate slides</h2>
-                <p style="margin-top: 20px;">Please try again.</p>
-            </div>
-            """
+            yield {
+                "html": f"""
+                <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
+                    <h2 style="color: #d9534f;">Failed to generate slides</h2>
+                    <p style="margin-top: 20px;">Please try again.</p>
+                </div>
+                """,
+                "files": []
+            }
             return
         
         audio_files = []
         validated_speaker_wav = await validate_and_convert_speaker_audio(speaker_audio)
         if not validated_speaker_wav:
             logger.error("Invalid speaker audio after conversion, skipping TTS")
-            yield f"""
-            <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
-                <h2 style="color: #d9534f;">Invalid speaker audio</h2>
-                <p style="margin-top: 20px;">Please upload a valid MP3 or WAV audio file and try again.</p>
-            </div>
-            """
+            yield {
+                "html": f"""
+                <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
+                    <h2 style="color: #d9534f;">Invalid speaker audio</h2>
+                    <p style="margin-top: 20px;">Please upload a valid MP3 or WAV audio file and try again.</p>
+                </div>
+                """,
+                "files": []
+            }
             return
         
         for i, script in enumerate(scripts):
@@ -668,8 +716,11 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
                 logger.error("Skipping audio for slide %d due to empty or invalid script", i + 1)
                 audio_files.append(None)
                 progress = 90 + ((i + 1) / len(scripts)) * 10
-                label = f"Generated audio for slide {i + 1}/{len(scripts)}..."
-                yield html_with_progress(label, progress)
+                label = f"Generating speech for slide {i + 1}/{len(scripts)}..."
+                yield {
+                    "html": html_with_progress(label, progress),
+                    "files": []
+                }
                 await asyncio.sleep(0.1)
                 continue
             
@@ -691,7 +742,10 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
                     audio_files.append(audio_file)
                     progress = 90 + ((i + 1) / len(scripts)) * 10
                     label = f"Generated audio for slide {i + 1}/{len(scripts)}..."
-                    yield html_with_progress(label, progress)
+                    yield {
+                        "html": html_with_progress(label, progress),
+                        "files": []
+                    }
                     await asyncio.sleep(0.1)
                     break
                 except Exception as e:
@@ -701,25 +755,25 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
                         audio_files.append(None)
                         progress = 90 + ((i + 1) / len(scripts)) * 10
                         label = f"Generated audio for slide {i + 1}/{len(scripts)}..."
-                        yield html_with_progress(label, progress)
+                        yield {
+                            "html": html_with_progress(label, progress),
+                            "files": []
+                        }
                         await asyncio.sleep(0.1)
                         break
         
         # Collect .txt files for download
         txt_files = [f for f in os.listdir(OUTPUT_DIR) if f.endswith('.txt')]
         txt_files.sort()  # Sort for consistent display
-        txt_links = ""
-        for txt_file in txt_files:
-            file_path = os.path.join(OUTPUT_DIR, txt_file)
-            txt_links += f'<a href="file/{file_path}" download>{txt_file}</a>  '
+        txt_file_paths = [os.path.join(OUTPUT_DIR, f) for f in txt_files]
         
-        # Generate audio timeline
+        # Generate audio timeline with playable audio elements
         audio_timeline = ""
         for i, audio_file in enumerate(audio_files):
             if audio_file:
-                audio_timeline += f'<span id="audio-{i+1}">{os.path.basename(audio_file)}</span>  '
+                audio_timeline += f'<audio id="audio-{i+1}" controls src="{audio_file}" style="margin: 0 5px;"></audio>'
             else:
-                audio_timeline += f'<span id="audio-{i+1}">slide_{i+1}.mp3</span>  '
+                audio_timeline += f'<span id="audio-{i+1}" style="margin: 0 5px;">slide_{i+1}.mp3 (not generated)</span>'
         
         slides_info = json.dumps({"slides": markdown_slides, "audioFiles": audio_files})
         
@@ -733,16 +787,12 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
                     {audio_timeline}
                 </div>
                 <div style="display: flex; justify-content: center; margin-bottom: 10px;">
-                    <button onclick="prevSlide()" style="border-radius: 50%; width: 40px; height: 40px; margin: 0 5px; font-size: 1.2em; cursor: pointer;">⏮</button>
-                    <button onclick="playAll()" style="border-radius: 50%; width: 40px; height: 40px; margin: 0 5px; font-size: 1.2em; cursor: pointer;">⏯</button>
-                    <button onclick="nextSlide()" style="border-radius: 50%; width: 40px; height: 40px; margin: 0 5px; font-size: 1.2em; cursor: pointer;">⏭</button>
+                    <button id="prev-btn" style="border-radius: 50%; width: 40px; height: 40px; margin: 0 5px; font-size: 1.2em; cursor: pointer;">⏮</button>
+                    <button id="play-btn" style="border-radius: 50%; width: 40px; height: 40px; margin: 0 5px; font-size: 1.2em; cursor: pointer;">⏯</button>
+                    <button id="next-btn" style="border-radius: 50%; width: 40px; height: 40px; margin: 0 5px; font-size: 1.2em; cursor: pointer;">⏭</button>
                     <button style="border-radius: 50%; width: 40px; height: 40px; margin: 0 5px; font-size: 1.2em; cursor: pointer;">☐</button>
                 </div>
             </div>
-            <div style="padding: 10px; text-align: center;">
-                <h4>Download Generated Files:</h4>
-                {txt_links}
-            </div>
         </div>
         <script>
             const lectureData = {slides_info};
@@ -752,12 +802,8 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
             let isPlayingAll = false;
 
             for (let i = 0; i < totalSlides; i++) {{
-                if (lectureData.audioFiles && lectureData.audioFiles[i]) {{
-                    const audio = new Audio(lectureData.audioFiles[i]);
-                    audioElements.push(audio);
-                }} else {{
-                    audioElements.push(null);
-                }}
+                const audio = document.getElementById(`audio-${{i+1}}`);
+                audioElements.push(audio);
             }}
 
             function renderSlide() {{
@@ -768,7 +814,7 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
             function updateSlide() {{
                 renderSlide();
                 audioElements.forEach(audio => {{
-                    if (audio) {{
+                    if (audio && audio.pause) {{
                         audio.pause();
                         audio.currentTime = 0;
                     }}
@@ -792,13 +838,15 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
             function playAll() {{
                 if (isPlayingAll) {{
                     audioElements.forEach(audio => {{
-                        if (audio) audio.pause();
+                        if (audio && audio.pause) audio.pause();
                     }});
                     isPlayingAll = false;
+                    document.getElementById('play-btn').innerText = '⏯';
                     return;
                 }}
 
                 isPlayingAll = true;
+                document.getElementById('play-btn').innerText = '⏸';
                 currentSlide = 0;
                 updateSlide();
                 playCurrentSlide();
@@ -807,11 +855,12 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
             function playCurrentSlide() {{
                 if (!isPlayingAll || currentSlide >= totalSlides) {{
                     isPlayingAll = false;
+                    document.getElementById('play-btn').innerText = '⏯';
                     return;
                 }}
 
                 const audio = audioElements[currentSlide];
-                if (audio) {{
+                if (audio && audio.play) {{
                     audio.play().then(() => {{
                         audio.addEventListener('ended', () => {{
                             currentSlide++;
@@ -820,6 +869,7 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
                                 playCurrentSlide();
                             }} else {{
                                 isPlayingAll = false;
+                                document.getElementById('play-btn').innerText = '⏯';
                             }}
                         }}, {{ once: true }});
                     }}).catch(e => {{
@@ -833,22 +883,33 @@ Example: 'Received {total_slides} slides and {total_slides} scripts. Lecture is
                 }}
             }}
 
+            // Attach event listeners
+            document.getElementById('prev-btn').addEventListener('click', prevSlide);
+            document.getElementById('play-btn').addEventListener('click', playAll);
+            document.getElementById('next-btn').addEventListener('click', nextSlide);
+
             // Initialize first slide
             renderSlide();
         </script>
         """
         logger.info("Lecture generation completed successfully")
-        yield html_output
+        yield {
+            "html": html_output,
+            "files": txt_file_paths
+        }
     
     except Exception as e:
         logger.error("Error during lecture generation: %s\n%s", str(e), traceback.format_exc())
-        yield f"""
-        <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
-            <h2 style="color: #d9534f;">Error during lecture generation</h2>
-            <p style="margin-top: 10px; font-size: 16px;">{str(e)}</p>
-            <p style="margin-top: 20px;">Please try again or adjust your inputs.</p>
-        </div>
-        """
+        yield {
+            "html": f"""
+            <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
+                <h2 style="color: #d9534f;">Error during lecture generation</h2>
+                <p style="margin-top: 10px; font-size: 16px;">{str(e)}</p>
+                <p style="margin-top: 20px;">Please try again or adjust your inputs.</p>
+            </div>
+            """,
+            "files": []
+        }
         return
 
 # Gradio interface
@@ -884,6 +945,7 @@ with gr.Blocks(title="Agent Feynman") as demo:
             </div>
             """
             slide_display = gr.HTML(label="Lecture Slides", value=default_slide_html)
+            file_output = gr.File(label="Download Generated Files")
     
     speaker_audio.change(
         fn=update_audio_preview,
@@ -894,7 +956,7 @@ with gr.Blocks(title="Agent Feynman") as demo:
     generate_btn.click(
         fn=on_generate,
         inputs=[api_service, api_key, serpapi_key, title, topic, instructions, lecture_type, speaker_audio, num_slides],
-        outputs=[slide_display]
+        outputs=[slide_display, file_output]
     )
 
 if __name__ == "__main__":