Update app.py

app.py

@@ -6,6 +6,8 @@ import asyncio
 import logging
 import torch
 import random
+import tempfile
+import zipfile
 from serpapi import GoogleSearch
 from pydantic import BaseModel
 from autogen_agentchat.agents import AssistantAgent
@@ -19,9 +21,10 @@ from autogen_ext.models.ollama import OllamaChatCompletionClient
 from markdown_pdf import MarkdownPdf, Section
 import traceback
 import soundfile as sf
-import tempfile
+import shutil
 from pydub import AudioSegment
 from TTS.api import TTS
+from gradio_pdf import PDF
 
 # Set up logging
 logging.basicConfig(
@@ -35,8 +38,6 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 # Set up environment
-OUTPUT_DIR = "outputs"
-os.makedirs(OUTPUT_DIR, exist_ok=True)
 os.environ["COQUI_TOS_AGREED"] = "1"
 
 # Define Pydantic model for slide data
@@ -82,17 +83,6 @@ def search_web(query: str, serpapi_key: str) -> str:
         logger.error("Unexpected error during search: %s", str(e))
         return f"Unexpected error during search: {str(e)}"
 
-# Define helper function for progress HTML
-def html_with_progress(label, progress):
-    return 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;">
-        <div style="width: 100%; background-color: #FFFFFF; border-radius: 10px; overflow: hidden; margin-bottom: 20px;">
-            <div style="width: {progress}%; height: 30px; background-color: #4CAF50; border-radius: 10px;"></div>
-        </div>
-        <h2 style="font-style: italic; color: #555;">{label}</h2>
-    </div>
-    """
-
 # Function to get model client based on selected service
 def get_model_client(service, api_key):
     if service == "OpenAI-gpt-4o-2024-08-06":
@@ -139,7 +129,7 @@ def clean_script_text(script):
     return script
 
 # Helper function to validate and convert speaker audio (MP3 or WAV)
-async def validate_and_convert_speaker_audio(speaker_audio):
+async def validate_and_convert_speaker_audio(speaker_audio, temp_dir):
     if not os.path.exists(speaker_audio):
         logger.error("Speaker audio file does not exist: %s", speaker_audio)
         return None
@@ -152,9 +142,8 @@ async def validate_and_convert_speaker_audio(speaker_audio):
             audio = AudioSegment.from_mp3(speaker_audio)
             # Convert to mono, 22050 Hz
             audio = audio.set_channels(1).set_frame_rate(22050)
-            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
-                audio.export(temp_file.name, format="wav")
-                speaker_wav = temp_file.name
+            speaker_wav = os.path.join(temp_dir, "speaker_converted.wav")
+            audio.export(speaker_wav, format="wav")
         elif ext == ".wav":
             speaker_wav = speaker_audio
         else:
@@ -172,9 +161,9 @@ async def validate_and_convert_speaker_audio(speaker_audio):
         if data.ndim == 2:
             logger.info("Converting stereo WAV to mono: %s", speaker_wav)
             data = data.mean(axis=1)
-            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
-                sf.write(temp_file.name, data, samplerate)
-                speaker_wav = temp_file.name
+            mono_wav = os.path.join(temp_dir, "speaker_mono.wav")
+            sf.write(mono_wav, data, samplerate)
+            speaker_wav = mono_wav
         
         logger.info("Validated speaker audio: %s", speaker_wav)
         return speaker_wav
@@ -280,7 +269,7 @@ def extract_json_from_message(message):
     return None
 
 # Function to generate Markdown and convert to PDF (portrait, centered)
-def generate_slides_pdf(slides):
+def generate_slides_pdf(slides, temp_dir):
     pdf = MarkdownPdf()
     
     for slide in slides:
@@ -299,12 +288,49 @@ def generate_slides_pdf(slides):
 """
         pdf.add_section(Section(markdown_content, toc=False))
     
-    pdf_file = os.path.join(OUTPUT_DIR, "slides.pdf")
+    pdf_file = os.path.join(temp_dir, "slides.pdf")
     pdf.save(pdf_file)
     
     logger.info("Generated PDF slides (portrait): %s", pdf_file)
     return pdf_file
 
+# Helper function to create ZIP file of outputs
+def create_outputs_zip(temp_dir, slides, audio_files, scripts):
+    zip_path = os.path.join(temp_dir, "lecture_outputs.zip")
+    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        # Add slides PDF
+        pdf_file = os.path.join(temp_dir, "slides.pdf")
+        if os.path.exists(pdf_file):
+            zipf.write(pdf_file, "slides.pdf")
+        
+        # Add audio files
+        for i, audio_file in enumerate(audio_files):
+            if audio_file and os.path.exists(audio_file):
+                zipf.write(audio_file, f"slide_{i+1}.wav")
+        
+        # Add raw and cleaned scripts
+        for i in range(len(slides)):
+            raw_script_file = os.path.join(temp_dir, f"slide_{i+1}_raw_script.txt")
+            cleaned_script_file = os.path.join(temp_dir, f"slide_{i+1}_script.txt")
+            if os.path.exists(raw_script_file):
+                zipf.write(raw_script_file, f"slide_{i+1}_raw_script.txt")
+            if os.path.exists(cleaned_script_file):
+                zipf.write(cleaned_script_file, f"slide_{i+1}_script.txt")
+    
+    logger.info("Created ZIP file: %s", zip_path)
+    return zip_path
+
+# Helper function for progress HTML
+def html_with_progress(label, progress):
+    return 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;">
+        <div style="width: 100%; background-color: #FFFFFF; border-radius: 10px; overflow: hidden; margin-bottom: 20px;">
+            <div style="width: {progress}%; height: 30px; background-color: #4CAF50; border-radius: 10px;"></div>
+        </div>
+        <h2 style="font-style: italic; color: #555;">{label}</h2>
+    </div>
+    """
+
 # Async function to update audio preview
 async def update_audio_preview(audio_file):
     if audio_file:
@@ -315,47 +341,38 @@ 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_with_progress("SerpApi key required. Please provide a valid key.", 0)
         return
 
-    # Initialize TTS model
-    tts = None
-    try:
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(device)
-        logger.info("TTS model initialized on %s", device)
-    except Exception as e:
-        logger.error("Failed to initialize TTS model: %s", str(e))
-        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;">Error: {str(e)}</p>
-            <p>Please ensure the Coqui TTS model is properly installed and try again.</p>
-        </div>
-        """
-        return
+    # Create temporary directory
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Initialize TTS model
+        tts = None
+        try:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(device)
+            logger.info("TTS model initialized on %s", device)
+        except Exception as e:
+            logger.error("Failed to initialize TTS model: %s", str(e))
+            yield html_with_progress(f"TTS model initialization failed: {str(e)}", 0)
+            return
 
-    model_client = get_model_client(api_service, api_key)
-    
-    research_agent = AssistantAgent(
-        name="research_agent",
-        model_client=model_client,
-        handoffs=["slide_agent"],
-        system_message="You are a Research Agent. Use the search_web tool to gather information on the topic and keywords from the initial message. Summarize the findings concisely in a single message, then use the handoff_to_slide_agent tool to pass the task to the Slide Agent. Do not produce any other output.",
-        tools=[search_web]
-    )
-    slide_agent = AssistantAgent(
-        name="slide_agent",
-        model_client=model_client,
-        handoffs=["script_agent"],
-        system_message=f"""
-You are a Slide Agent. Using the research from the conversation history and the specified number of content slides ({num_slides}), generate exactly {num_slides} content slides, plus one quiz slide, one assignment slide, and one thank-you slide, for a total of {num_slides + 3} slides. Output ONLY a JSON array wrapped in ```json ... ``` in a TextMessage, where each slide is an object with 'title' and 'content' keys. Do not include any explanatory text, comments, or other messages. Ensure the JSON is valid and contains exactly {num_slides + 3} slides before proceeding. After outputting the JSON, use the handoff_to_script_agent tool to pass the task to the Script Agent.
-Example output for 2 content slides:
+        model_client = get_model_client(api_service, api_key)
+        
+        research_agent = AssistantAgent(
+            name="research_agent",
+            model_client=model_client,
+            handoffs=["slide_agent"],
+            system_message="You are a Research Agent. Use the search_web tool to gather information on the topic and keywords from the initial message. Summarize the findings concisely in a single message, then use the handoff_to_slide_agent tool to pass the task to the Slide Agent. Do not produce any other output.",
+            tools=[search_web]
+        )
+        slide_agent = AssistantAgent(
+            name="slide_agent",
+            model_client=model_client,
+            handoffs=["script_agent"],
+            system_message=f"""
+You are a Slide Agent. Using the research from the conversation history, generate EXACTLY {num_slides} content slides on the topic, plus 1 quiz slide, 1 assignment slide, and 1 thank-you slide, for a TOTAL of {num_slides + 3} slides. Output ONLY a JSON array wrapped in ```json ... ``` in a TextMessage, with each slide as an object with 'title' and 'content' keys. Ensure the JSON is valid and contains precisely {num_slides + 3} slides. If the slide count is incorrect, adjust the output to meet this requirement before proceeding. Do not include explanatory text or comments. After outputting the JSON, use the handoff_to_script_agent tool.
+Example for 2 content slides:
 ```json
 [
     {{"title": "Slide 1", "content": "Content for slide 1"}},
@@ -365,14 +382,14 @@ Example output for 2 content slides:
     {{"title": "Thank You", "content": "Thank you message"}}
 ]
 ```""",
-        output_content_type=None,
-        reflect_on_tool_use=False
-    )
-    script_agent = AssistantAgent(
-        name="script_agent",
-        model_client=model_client,
-        handoffs=["feynman_agent"],
-        system_message=f"""
+            output_content_type=None,
+            reflect_on_tool_use=False
+        )
+        script_agent = AssistantAgent(
+            name="script_agent",
+            model_client=model_client,
+            handoffs=["feynman_agent"],
+            system_message=f"""
 You are a Script Agent. Access the JSON array of {num_slides + 3} slides from the conversation history. Generate a narration script (1-2 sentences) for each of the {num_slides + 3} slides, summarizing its content in a natural, conversational tone as a speaker would, including occasional non-verbal words (e.g., "um," "you know," "like"). Output ONLY a JSON array wrapped in ```json ... ``` with exactly {num_slides + 3} strings, one script per slide, in the same order. Ensure the JSON is valid and complete. After outputting, use the handoff_to_feynman_agent tool. If scripts cannot be generated, retry once.
 Example for 1 content slide:
 ```json
@@ -383,392 +400,373 @@ Example for 1 content slide:
     "Thanks for, um, attending today!"
 ]
 ```""",
-        output_content_type=None,
-        reflect_on_tool_use=False
-    )
-    feynman_agent = AssistantAgent(
-        name="feynman_agent",
-        model_client=model_client,
-        handoffs=[],
-        system_message=f"""
+            output_content_type=None,
+            reflect_on_tool_use=False
+        )
+        feynman_agent = AssistantAgent(
+            name="feynman_agent",
+            model_client=model_client,
+            handoffs=[],
+            system_message=f"""
 You are Agent Feynman. Review the slides and scripts from the conversation history to ensure coherence, completeness, and that exactly {num_slides + 3} slides and {num_slides + 3} scripts are received. Output a confirmation message summarizing the number of slides and scripts received. If slides or scripts are missing, invalid, or do not match the expected count ({num_slides + 3}), report the issue clearly. Use 'TERMINATE' to signal completion.
 Example: 'Received {num_slides + 3} slides and {num_slides + 3} scripts. Lecture is coherent. TERMINATE'
 """)
-    
-    swarm = Swarm(
-        participants=[research_agent, slide_agent, script_agent, feynman_agent],
-        termination_condition=HandoffTermination(target="user") | TextMentionTermination("TERMINATE")
-    )
-    
-    progress = 0
-    label = "Research: in progress..."
-    yield html_with_progress(label, progress)
-    await asyncio.sleep(0.1)
-    
-    initial_message = f"""
-    Lecture Title: {title}
-    Topic: {topic}
-    Additional Instructions: {instructions}
-    Audience: {lecture_type}
-    Number of Content Slides: {num_slides}
-    Please start by researching the topic.
-    """
-    logger.info("Starting lecture generation for topic: %s", topic)
-    
-    slides = None
-    scripts = None
-    error_html = """
-    <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 lecture materials</h2>
-        <p style="margin-top: 20px;">Please try again with different parameters or a different model.</p>
-    </div>
-    """
-    
-    try:
-        logger.info("Research Agent starting...")
-        task_result = await Console(swarm.run_stream(task=initial_message))
-        logger.info("Swarm execution completed")
         
-        script_retry_count = 0
-        max_script_retries = 2
+        swarm = Swarm(
+            participants=[research_agent, slide_agent, script_agent, feynman_agent],
+            termination_condition=HandoffTermination(target="user") | TextMentionTermination("TERMINATE")
+        )
         
-        for message in task_result.messages:
-            source = getattr(message, 'source', getattr(message, 'sender', None))
-            logger.debug("Processing message from %s, type: %s, content: %s", source, type(message), message.to_text() if hasattr(message, 'to_text') else str(message))
-            
-            if isinstance(message, HandoffMessage):
-                logger.info("Handoff from %s to %s, Context: %s", source, message.target, message.context)
-                if source == "research_agent" and message.target == "slide_agent":
-                    progress = 25
-                    label = "Slides: generating..."
-                    yield html_with_progress(label, progress)
-                    await asyncio.sleep(0.1)
-                elif source == "slide_agent" and message.target == "script_agent":
-                    if slides is None:
-                        logger.warning("Slide Agent handoff without slides JSON")
+        progress = 0
+        label = "Research: in progress..."
+        yield html_with_progress(label, progress)
+        await asyncio.sleep(0.1)
+        
+        initial_message = f"""
+        Lecture Title: {title}
+        Topic: {topic}
+        Additional Instructions: {instructions}
+        Audience: {lecture_type}
+        Number of Content Slides: {num_slides}
+        Please start by researching the topic.
+        """
+        logger.info("Starting lecture generation for topic: %s", topic)
+        
+        slides = None
+        scripts = None
+        max_slide_retries = 2
+        slide_retry_count = 0
+        
+        while slide_retry_count <= max_slide_retries:
+            try:
+                logger.info("Research Agent starting (Slide attempt %d/%d)", slide_retry_count + 1, max_slide_retries)
+                task_result = await Console(swarm.run_stream(task=initial_message))
+                logger.info("Swarm execution completed")
+                
+                script_retry_count = 0
+                max_script_retries = 2
+                
+                for message in task_result.messages:
+                    source = getattr(message, 'source', getattr(message, 'sender', None))
+                    logger.debug("Processing message from %s, type: %s, content: %s", source, type(message), message.to_text() if hasattr(message, 'to_text') else str(message))
+                    
+                    if isinstance(message, HandoffMessage):
+                        logger.info("Handoff from %s to %s, Context: %s", source, message.target, message.context)
+                        if source == "research_agent" and message.target == "slide_agent":
+                            progress = 25
+                            label = "Slides: generating..."
+                            yield html_with_progress(label, progress)
+                            await asyncio.sleep(0.1)
+                        elif source == "slide_agent" and message.target == "script_agent":
+                            if slides is None:
+                                logger.warning("Slide Agent handoff without slides JSON")
+                                extracted_json = extract_json_from_message(message)
+                                if extracted_json:
+                                    slides = extracted_json
+                                    logger.info("Extracted slides JSON from HandoffMessage context: %s", slides)
+                            if slides is None:
+                                label = "Slides: failed to generate..."
+                                yield html_with_progress(label, progress)
+                                await asyncio.sleep(0.1)
+                            progress = 50
+                            label = "Scripts: generating..."
+                            yield html_with_progress(label, progress)
+                            await asyncio.sleep(0.1)
+                        elif source == "script_agent" and message.target == "feynman_agent":
+                            if scripts is None:
+                                logger.warning("Script Agent handoff without scripts JSON")
+                                extracted_json = extract_json_from_message(message)
+                                if extracted_json:
+                                    scripts = extracted_json
+                                    logger.info("Extracted scripts JSON from HandoffMessage context: %s", scripts)
+                            progress = 75
+                            label = "Review: in progress..."
+                            yield html_with_progress(label, progress)
+                            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)
+                        await asyncio.sleep(0.1)
+                    
+                    elif source == "slide_agent" and isinstance(message, (TextMessage, StructuredMessage)):
+                        logger.debug("Slide Agent message received: %s", message.to_text())
                         extracted_json = extract_json_from_message(message)
                         if extracted_json:
                             slides = extracted_json
-                            logger.info("Extracted slides JSON from HandoffMessage context: %s", slides)
-                    if slides is None:
-                        label = "Slides: failed to generate..."
-                        yield html_with_progress(label, progress)
-                        await asyncio.sleep(0.1)
-                    progress = 50
-                    label = "Scripts: generating..."
-                    yield html_with_progress(label, progress)
-                    await asyncio.sleep(0.1)
-                elif source == "script_agent" and message.target == "feynman_agent":
-                    if scripts is None:
-                        logger.warning("Script Agent handoff without scripts JSON")
+                            logger.info("Slide Agent generated %d slides: %s", len(slides), slides)
+                            expected_slide_count = num_slides + 3
+                            if len(slides) != expected_slide_count:
+                                logger.warning("Generated %d slides, expected %d. Retrying...", len(slides), expected_slide_count)
+                                slide_retry_count += 1
+                                if slide_retry_count <= max_slide_retries:
+                                    # Re-prompt slide agent
+                                    retry_message = TextMessage(
+                                        content=f"Please generate EXACTLY {num_slides} content slides plus 1 quiz, 1 assignment, and 1 thank-you slide (total {num_slides + 3}).",
+                                        source="user",
+                                        recipient="slide_agent"
+                                    )
+                                    task_result.messages.append(retry_message)
+                                    slides = None
+                                    continue
+                                else:
+                                    yield html_with_progress(f"Failed to generate correct number of slides after {max_slide_retries} retries. Expected {expected_slide_count}, got {len(slides)}.", progress)
+                                    return
+                            # Save slide content to individual files
+                            for i, slide in enumerate(slides):
+                                content_file = os.path.join(temp_dir, f"slide_{i+1}_content.txt")
+                                try:
+                                    with open(content_file, "w", encoding="utf-8") as f:
+                                        f.write(slide["content"])
+                                    logger.info("Saved slide content to %s: %s", content_file, slide["content"])
+                                except Exception as e:
+                                    logger.error("Error saving slide content to %s: %s", content_file, str(e))
+                            progress = 50
+                            label = "Scripts: generating..."
+                            yield html_with_progress(label, progress)
+                            await asyncio.sleep(0.1)
+                        else:
+                            logger.warning("No JSON extracted from slide_agent message: %s", message.to_text())
+                    
+                    elif source == "script_agent" and isinstance(message, (TextMessage, StructuredMessage)):
+                        logger.debug("Script Agent message received: %s", message.to_text())
                         extracted_json = extract_json_from_message(message)
                         if extracted_json:
                             scripts = extracted_json
-                            logger.info("Extracted scripts JSON from HandoffMessage context: %s", scripts)
-                    progress = 75
-                    label = "Review: in progress..."
-                    yield html_with_progress(label, progress)
-                    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)
-                await asyncio.sleep(0.1)
-            
-            elif source == "slide_agent" and isinstance(message, (TextMessage, StructuredMessage)):
-                logger.debug("Slide Agent message received: %s", message.to_text())
-                extracted_json = extract_json_from_message(message)
-                if extracted_json:
-                    slides = extracted_json
-                    logger.info("Slide Agent generated %d slides: %s", len(slides), slides)
-                    # Save slide content to individual files
-                    for i, slide in enumerate(slides):
-                        content_file = os.path.join(OUTPUT_DIR, f"slide_{i+1}_content.txt")
-                        try:
-                            with open(content_file, "w", encoding="utf-8") as f:
-                                f.write(slide["content"])
-                            logger.info("Saved slide content to %s: %s", content_file, slide["content"])
-                        except Exception as e:
-                            logger.error("Error saving slide content to %s: %s", content_file, str(e))
-                    progress = 50
-                    label = "Scripts: generating..."
-                    yield html_with_progress(label, progress)
-                    await asyncio.sleep(0.1)
-                else:
-                    logger.warning("No JSON extracted from slide_agent message: %s", message.to_text())
-            
-            elif source == "script_agent" and isinstance(message, (TextMessage, StructuredMessage)):
-                logger.debug("Script Agent message received: %s", message.to_text())
-                extracted_json = extract_json_from_message(message)
-                if extracted_json:
-                    scripts = extracted_json
-                    logger.info("Script Agent generated scripts for %d slides: %s", len(scripts), scripts)
-                    # Save raw scripts to individual files
-                    for i, script in enumerate(scripts):
-                        script_file = os.path.join(OUTPUT_DIR, f"slide_{i+1}_raw_script.txt")
-                        try:
-                            with open(script_file, "w", encoding="utf-8") as f:
-                                f.write(script)
-                            logger.info("Saved raw script to %s: %s", script_file, script)
-                        except Exception as e:
-                            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)
-                    await asyncio.sleep(0.1)
-                else:
-                    logger.warning("No JSON extracted from script_agent message: %s", message.to_text())
-                    if script_retry_count < max_script_retries:
-                        script_retry_count += 1
-                        logger.info("Retrying script generation (attempt %d/%d)", script_retry_count, max_script_retries)
-                        # Re-prompt script agent
-                        retry_message = TextMessage(
-                            content="Please generate scripts for the slides as per your instructions.",
-                            source="user",
-                            recipient="script_agent"
-                        )
-                        task_result.messages.append(retry_message)
-                        continue
-            
-            elif source == "feynman_agent" and isinstance(message, TextMessage) and "TERMINATE" in message.content:
-                logger.info("Feynman Agent completed lecture review: %s", message.content)
-                progress = 90  # Set to 90% before audio generation
-                label = "Lecture materials ready. Generating audio..."
-                yield html_with_progress(label, progress)
-                await asyncio.sleep(0.1)
-        
-        logger.info("Slides state: %s", "Generated" if slides else "None")
-        logger.info("Scripts state: %s", "Generated" if scripts else "None")
-        if not slides or not scripts:
-            error_message = f"Failed to generate {'slides and scripts' if not slides and not scripts else 'slides' if not slides else 'scripts'}"
-            error_message += f". Received {len(slides) if slides else 0} slides and {len(scripts) if scripts else 0} scripts."
-            logger.error("%s", error_message)
-            logger.debug("Dumping all messages for debugging:")
-            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>
-            """
-            return
-        
-        expected_slide_count = num_slides + 3
-        if len(slides) != expected_slide_count:
-            logger.error("Expected %d slides (including %d content slides + 3), but received %d", expected_slide_count, num_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 {expected_slide_count} slides ({num_slides} content slides + quiz, assignment, thank-you), but generated {len(slides)}. Please try again.</p>
-            </div>
-            """
-            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>
-            """
-            return
-        
-        if len(scripts) != expected_slide_count:
-            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>
-            """
-            return
-        
-        # Generate PDF from slides
-        pdf_file = generate_slides_pdf(slides)
-        pdf_path = f"file://{os.path.abspath(pdf_file)}"
-        
-        audio_files = []
-        speaker_audio = speaker_audio if speaker_audio else "feynman.mp3"
-        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>
-            """
-            return
-        
-        # Process audio generation sequentially with retries
-        for i, script in enumerate(scripts):
-            cleaned_script = clean_script_text(script)
-            audio_file = os.path.join(OUTPUT_DIR, f"slide_{i+1}.wav")
-            script_file = os.path.join(OUTPUT_DIR, f"slide_{i+1}_script.txt")
-            
-            # Save cleaned script
-            try:
-                with open(script_file, "w", encoding="utf-8") as f:
-                    f.write(cleaned_script or "")
-                logger.info("Saved cleaned script to %s: %s", script_file, cleaned_script)
-            except Exception as e:
-                logger.error("Error saving cleaned script to %s: %s", script_file, str(e))
-            
-            if not cleaned_script:
-                logger.error("Skipping audio for slide %d due to empty or invalid script", i + 1)
-                audio_files.append(None)
-                # Update progress (even for skipped slides)
-                progress = 90 + ((i + 1) / len(scripts)) * 10  # Progress from 90% to 100%
-                label = f"Generated audio for slide {i + 1}/{len(scripts)}..."
-                yield html_with_progress(label, progress)
-                await asyncio.sleep(0.1)
-                continue
-            
-            max_retries = 2
-            for attempt in range(max_retries + 1):
-                try:
-                    current_text = cleaned_script
-                    if attempt > 0:
-                        sentences = re.split(r"[.!?]+", cleaned_script)
-                        sentences = [s.strip() for s in sentences if s.strip()][:2]
-                        current_text = ". ".join(sentences) + "."
-                        logger.info("Retry %d for slide %d with simplified text: %s", attempt, i + 1, current_text)
+                            logger.info("Script Agent generated scripts for %d slides: %s", len(scripts), scripts)
+                            # Save raw scripts to individual files
+                            for i, script in enumerate(scripts):
+                                script_file = os.path.join(temp_dir, f"slide_{i+1}_raw_script.txt")
+                                try:
+                                    with open(script_file, "w", encoding="utf-8") as f:
+                                        f.write(script)
+                                    logger.info("Saved raw script to %s: %s", script_file, script)
+                                except Exception as e:
+                                    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)
+                            await asyncio.sleep(0.1)
+                        else:
+                            logger.warning("No JSON extracted from script_agent message: %s", message.to_text())
+                            if script_retry_count < max_script_retries:
+                                script_retry_count += 1
+                                logger.info("Retrying script generation (attempt %d/%d)", script_retry_count, max_script_retries)
+                                # Re-prompt script agent
+                                retry_message = TextMessage(
+                                    content="Please generate scripts for the slides as per your instructions.",
+                                    source="user",
+                                    recipient="script_agent"
+                                )
+                                task_result.messages.append(retry_message)
+                                continue
                     
-                    success = generate_xtts_audio(tts, current_text, validated_speaker_wav, audio_file)
-                    if not success:
-                        raise RuntimeError("TTS generation failed")
+                    elif source == "feynman_agent" and isinstance(message, TextMessage) and "TERMINATE" in message.content:
+                        logger.info("Feynman Agent completed lecture review: %s", message.content)
+                        progress = 90
+                        label = "Lecture materials ready. Generating audio..."
+                        yield html_with_progress(label, progress)
+                        await asyncio.sleep(0.1)
+                
+                logger.info("Slides state: %s", "Generated" if slides else "None")
+                logger.info("Scripts state: %s", "Generated" if scripts else "None")
+                if not slides or not scripts:
+                    error_message = f"Failed to generate {'slides and scripts' if not slides and not scripts else 'slides' if not slides else 'scripts'}"
+                    error_message += f". Received {len(slides) if slides else 0} slides and {len(scripts) if scripts else 0} scripts."
+                    logger.error("%s", error_message)
+                    yield html_with_progress(error_message, progress)
+                    return
+                
+                expected_slide_count = num_slides + 3
+                if len(slides) != expected_slide_count:
+                    logger.error("Final validation failed: Expected %d slides, received %d", expected_slide_count, len(slides))
+                    yield html_with_progress(f"Incorrect number of slides. Expected {expected_slide_count}, got {len(slides)}.", progress)
+                    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 html_with_progress("Invalid script format. Scripts must be a list of strings.", progress)
+                    return
+                
+                if len(scripts) != expected_slide_count:
+                    logger.error("Mismatch between number of slides (%d) and scripts (%d)", len(slides), len(scripts))
+                    yield html_with_progress(f"Mismatch in slides and scripts. Generated {len(slides)} slides but {len(scripts)} scripts.", progress)
+                    return
+                
+                # Generate PDF from slides
+                pdf_file = generate_slides_pdf(slides, temp_dir)
+                
+                audio_files = []
+                speaker_audio = speaker_audio if speaker_audio else "feynman.mp3"
+                validated_speaker_wav = await validate_and_convert_speaker_audio(speaker_audio, temp_dir)
+                if not validated_speaker_wav:
+                    logger.error("Invalid speaker audio after conversion, skipping TTS")
+                    yield html_with_progress("Invalid speaker audio. Please upload a valid MP3 or WAV file.", progress)
+                    return
+                
+                # Process audio generation sequentially with retries
+                for i, script in enumerate(scripts):
+                    cleaned_script = clean_script_text(script)
+                    audio_file = os.path.join(temp_dir, f"slide_{i+1}.wav")
+                    script_file = os.path.join(temp_dir, f"slide_{i+1}_script.txt")
+                    
+                    # Save cleaned script
+                    try:
+                        with open(script_file, "w", encoding="utf-8") as f:
+                            f.write(cleaned_script or "")
+                        logger.info("Saved cleaned script to %s: %s", script_file, cleaned_script)
+                    except Exception as e:
+                        logger.error("Error saving cleaned script to %s: %s", script_file, str(e))
                     
-                    logger.info("Generated audio for slide %d: %s", i + 1, audio_file)
-                    audio_files.append(audio_file)
-                    progress = 90 + ((i + 1) / len(scripts)) * 10  # Progress from 90% to 100%
-                    label = f"Generated audio for slide {i + 1}/{len(scripts)}..."
-                    yield html_with_progress(label, progress)
-                    await asyncio.sleep(0.1)
-                    break
-                except Exception as e:
-                    logger.error("Error generating audio for slide %d (attempt %d): %s\n%s", i + 1, attempt, str(e), traceback.format_exc())
-                    if attempt == max_retries:
-                        logger.error("Max retries reached for slide %d, skipping", i + 1)
+                    if not cleaned_script:
+                        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  # Progress from 90% to 100%
+                        progress = 90 + ((i + 1) / len(scripts)) * 10
                         label = f"Generated audio for slide {i + 1}/{len(scripts)}..."
                         yield html_with_progress(label, progress)
                         await asyncio.sleep(0.1)
-                        break
-        
-        audio_files = [f"file://{os.path.abspath(f)}" if f else None for f in audio_files]
-        
-        slides_info = json.dumps({"slides": [
-            {"title": slide["title"], "content": slide["content"]}
-            for slide in slides
-        ], "audioFiles": audio_files})
-        
-        html_output = f"""
-        <div id="lecture-container" style="height: 700px; border: 1px solid #ddd; border-radius: 8px; display: flex; flex-direction: column; justify-content: space-between;">
-            <div id="slide-content" style="flex: 1; overflow: hidden;">
-                <iframe id="pdf-viewer" src="https://mozilla.github.io/pdf.js/web/viewer.html?file={pdf_path}" style="width: 100%; height: 100%; border: none;"></iframe>
-            </div>
-            <div style="padding: 20px;">
-                <div id="progress-bar" style="width: 100%; height: 5px; background-color: #ddd; border-radius: 2px; margin-bottom: 10px;">
-                    <div id="progress-fill" style="width: {(1/len(slides)*100)}%; height: 100%; background-color: #4CAF50; border-radius: 2px;"></div>
-                </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="togglePlay()" 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>
+                        continue
+                    
+                    max_retries = 2
+                    for attempt in range(max_retries + 1):
+                        try:
+                            current_text = cleaned_script
+                            if attempt > 0:
+                                sentences = re.split(r"[.!?]+", cleaned_script)
+                                sentences = [s.strip() for s in sentences if s.strip()][:2]
+                                current_text = ". ".join(sentences) + "."
+                                logger.info("Retry %d for slide %d with simplified text: %s", attempt, i + 1, current_text)
+                            
+                            success = generate_xtts_audio(tts, current_text, validated_speaker_wav, audio_file)
+                            if not success:
+                                raise RuntimeError("TTS generation failed")
+                            
+                            logger.info("Generated audio for slide %d: %s", i + 1, audio_file)
+                            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)
+                            await asyncio.sleep(0.1)
+                            break
+                        except Exception as e:
+                            logger.error("Error generating audio for slide %d (attempt %d): %s\n%s", i + 1, attempt, str(e), traceback.format_exc())
+                            if attempt == max_retries:
+                                logger.error("Max retries reached for slide %d, skipping", 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)
+                                await asyncio.sleep(0.1)
+                                break
+                
+                # Create ZIP file of all outputs
+                zip_path = create_outputs_zip(temp_dir, slides, audio_files, scripts)
+                
+                # Prepare UI output
+                slides_info = json.dumps({"slides": [
+                    {"title": slide["title"], "content": slide["content"]}
+                    for slide in slides
+                ], "audioFiles": audio_files})
+                
+                html_output = f"""
+                <div id="lecture-container" style="height: 700px; border: 1px solid #ddd; border-radius: 8px; display: flex; flex-direction: column; justify-content: space-between;">
+                    <div id="slide-content" style="flex: 1; overflow: auto;">
+                        <div id="pdf-viewer"></div>
+                    </div>
+                    <div style="padding: 20px;">
+                        <div id="progress-bar" style="width: 100%; height: 5px; background-color: #ddd; border-radius: 2px; margin-bottom: 10px;">
+                            <div id="progress-fill" style="width: {(1/len(slides)*100)}%; height: 100%; background-color: #4CAF50; border-radius: 2px;"></div>
+                        </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="togglePlay()" 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>
+                        </div>
+                        <p id="slide-counter" style="text-align: center;">Slide 1 of {len(slides)}</p>
+                    </div>
                 </div>
-                <p id="slide-counter" style="text-align: center;">Slide 1 of {len(slides)}</p>
-            </div>
-        </div>
-        <script>
-            const lectureData = {slides_info};
-            let currentSlide = 0;
-            const totalSlides = lectureData.slides.length;
-            const slideCounter = document.getElementById('slide-counter');
-            const progressFill = document.getElementById('progress-fill');
-            let audioElements = [];
-            let currentAudio = null;
-            const pdfViewer = document.getElementById('pdf-viewer');
+                <script>
+                    const lectureData = {slides_info};
+                    let currentSlide = 0;
+                    const totalSlides = lectureData.slides.length;
+                    const slideCounter = document.getElementById('slide-counter');
+                    const progressFill = document.getElementById('progress-fill');
+                    let audioElements = [];
+                    let currentAudio = null;
 
-            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);
-                }}
-            }}
+                    for (let i = 0; i < totalSlides; i++) {{
+                        if (lectureData.audioFiles && lectureData.audioFiles[i]) {{
+                            const audio = new Audio('file://' + lectureData.audioFiles[i]);
+                            audioElements.push(audio);
+                        }} else {{
+                            audioElements.push(null);
+                        }}
+                    }}
 
-            function updateSlide() {{
-                pdfViewer.src = `https://mozilla.github.io/pdf.js/web/viewer.html?file={pdf_path}#page=${{currentSlide + 1}}`;
-                slideCounter.textContent = `Slide ${{currentSlide + 1}} of ${{totalSlides}}`;
-                progressFill.style.width = `${{(currentSlide + 1) / totalSlides * 100}}%`;
+                    function updateSlide() {{
+                        slideCounter.textContent = `Slide ${{currentSlide + 1}} of ${{totalSlides}}`;
+                        progressFill.style.width = `${{(currentSlide + 1) / totalSlides * 100}}%`;
 
-                if (currentAudio) {{
-                    currentAudio.pause();
-                    currentAudio.currentTime = 0;
-                }}
+                        if (currentAudio) {{
+                            currentAudio.pause();
+                            currentAudio.currentTime = 0;
+                        }}
 
-                if (audioElements[currentSlide]) {{
-                    currentAudio = audioElements[currentSlide];
-                    currentAudio.play().catch(e => console.error('Audio play failed:', e));
-                }} else {{
-                    currentAudio = null;
-                }}
-            }}
+                        if (audioElements[currentSlide]) {{
+                            currentAudio = audioElements[currentSlide];
+                            currentAudio.play().catch(e => console.error('Audio play failed:', e));
+                        }} else {{
+                            currentAudio = null;
+                        }}
+                    }}
 
-            function prevSlide() {{
-                if (currentSlide > 0) {{
-                    currentSlide--;
-                    updateSlide();
-                }}
-            }}
+                    function prevSlide() {{
+                        if (currentSlide > 0) {{
+                            currentSlide--;
+                            updateSlide();
+                        }}
+                    }}
 
-            function nextSlide() {{
-                if (currentSlide < totalSlides - 1) {{
-                    currentSlide++;
-                    updateSlide();
-                }}
-            }}
+                    function nextSlide() {{
+                        if (currentSlide < totalSlides - 1) {{
+                            currentSlide++;
+                            updateSlide();
+                        }}
+                    }}
 
-            function togglePlay() {{
-                if (!audioElements[currentSlide]) return;
-                if (currentAudio.paused) {{
-                    currentAudio.play().catch(e => console.error('Audio play failed:', e));
-                }} else {{
-                    currentAudio.pause();
-                }}
-            }}
+                    function togglePlay() {{
+                        if (!audioElements[currentSlide]) return;
+                        if (currentAudio.paused) {{
+                            currentAudio.play().catch(e => console.error('Audio play failed:', e));
+                        }} else {{
+                            currentAudio.pause();
+                        }}
+                    }}
 
-            audioElements.forEach((audio, index) => {{
-                if (audio) {{
-                    audio.addEventListener('ended', () => {{
-                        if (index < totalSlides - 1) {{
-                            nextSlide();
+                    audioElements.forEach((audio, index) => {{
+                        if (audio) {{
+                            audio.addEventListener('ended', () => {{
+                                if (index < totalSlides - 1) {{
+                                    nextSlide();
+                                }}
+                            }});
                         }}
                     }});
-                }}
-            }});
-        </script>
-        """
-        logger.info("Lecture generation completed successfully")
-        yield html_output
-    
-    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>
-        """
-        return
+                </script>
+                """
+                yield {
+                    "pdf": pdf_file,
+                    "html": html_output,
+                    "zip": zip_path
+                }
+                return
+            
+            except Exception as e:
+                logger.error("Error during lecture generation: %s\n%s", str(e), traceback.format_exc())
+                yield html_with_progress(f"Error during lecture generation: {str(e)}", progress)
+                return
 
 # Gradio interface
 with gr.Blocks(title="Agent Feynman") as demo:
@@ -803,6 +801,8 @@ with gr.Blocks(title="Agent Feynman") as demo:
             </div>
             """
             slide_display = gr.HTML(label="Lecture Slides", value=default_slide_html)
+            pdf_display = gr.PDF(label="Lecture Slides PDF")
+            outputs_zip = gr.File(label="Download Outputs (PDF, Audio, Scripts)")
     
     speaker_audio.change(
         fn=update_audio_preview,
@@ -813,7 +813,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, pdf_display, outputs_zip]
     )
 
 if __name__ == "__main__":