Update app.py

app.py

@@ -16,19 +16,15 @@ class ConsoleOutput:
     def get_output(self):
         return "\n".join(self.messages)
 
-def process_pdf(pdf_file, progress=gr.Progress()):
+def process_screenplay(pdf_file, progress=gr.Progress()):
     if pdf_file is None:
         raise gr.Error("Please upload a PDF file")
     
     console = ConsoleOutput()
-    
-    # Set up logging to console
-    logging.basicConfig(level=logging.DEBUG,
-                       format='%(asctime)s - %(levelname)s - %(message)s')
-    
+    logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
     logger = logging.getLogger()
-    logger.handlers = []  # Remove existing handlers
-    
+    logger.handlers = []
+
     class GradioHandler(logging.Handler):
         def emit(self, record):
             msg = self.format(record)
@@ -38,21 +34,43 @@ def process_pdf(pdf_file, progress=gr.Progress()):
     
     try:
         processor = GeminiProcessor()
-        coverage_gen = CoverageGenerator()
-        
-        # Clean screenplay
-        progress(0.2, desc="Processing screenplay...")
+        progress(0.5, desc="Processing screenplay...")
         cleaned_path = Path("cleaned_screenplay_long.txt")
         success = processor.process_screenplay(pdf_file.name, str(cleaned_path))
+        
         if not success:
             raise gr.Error("Failed to process screenplay")
         
         with open(cleaned_path, 'r') as f:
             cleaned_text = f.read()
-        
-        # Generate coverage
-        progress(0.6, desc="Generating coverage...")
+            
+        progress(1.0, desc="Complete!")
+        return cleaned_text, gr.Button.update(interactive=True), console.get_output()
+    
+    except Exception as e:
+        error_msg = f"Error: {str(e)}"
+        console.write(error_msg)
+        raise gr.Error(error_msg)
+
+def generate_coverage(progress=gr.Progress()):
+    console = ConsoleOutput()
+    logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+    logger = logging.getLogger()
+    logger.handlers = []
+
+    class GradioHandler(logging.Handler):
+        def emit(self, record):
+            msg = self.format(record)
+            console.write(msg)
+            
+    logger.addHandler(GradioHandler())
+    
+    try:
+        coverage_gen = CoverageGenerator()
+        progress(0.5, desc="Generating coverage...")
+        cleaned_path = Path("cleaned_screenplay_long.txt")
         success = coverage_gen.generate_coverage(cleaned_path)
+        
         if not success:
             raise gr.Error("Failed to generate coverage")
             
@@ -60,7 +78,7 @@ def process_pdf(pdf_file, progress=gr.Progress()):
             coverage = f.read()
             
         progress(1.0, desc="Complete!")
-        return cleaned_text, coverage, console.get_output()
+        return coverage, console.get_output()
         
     except Exception as e:
         error_msg = f"Error: {str(e)}"
@@ -75,37 +93,26 @@ with gr.Blocks(title="Screenplay Coverage Generator") as demo:
     
     with gr.Row():
         process_btn = gr.Button("Process Screenplay")
+        coverage_btn = gr.Button("Generate Coverage", interactive=False)
     
     with gr.Row():
-        console = gr.Textbox(
-            label="Console Output",
-            lines=10,
-            max_lines=30,
-            autoscroll=True,
-            container=True,
-            show_copy_button=True
-        )
+        console = gr.Textbox(label="Console Output", lines=10, max_lines=30, autoscroll=True, show_copy_button=True)
     
     with gr.Tabs():
         with gr.TabItem("Cleaned Screenplay"):
-            cleaned_output = gr.Textbox(
-                label="Cleaned Screenplay",
-                lines=10,
-                max_lines=30,
-                show_copy_button=True
-            )
+            cleaned_output = gr.Textbox(label="Cleaned Screenplay", lines=10, show_copy_button=True)
         with gr.TabItem("Coverage"):
-            coverage_output = gr.Textbox(
-                label="Coverage Document",
-                lines=10,
-                max_lines=30,
-                show_copy_button=True
-            )
+            coverage_output = gr.Textbox(label="Coverage Document", lines=10, show_copy_button=True)
     
     process_btn.click(
-        fn=process_pdf,
+        fn=process_screenplay,
         inputs=[file_input],
-        outputs=[cleaned_output, coverage_output, console]
+        outputs=[cleaned_output, coverage_btn, console]
+    )
+    
+    coverage_btn.click(
+        fn=generate_coverage,
+        outputs=[coverage_output, console]
     )
 
 if __name__ == "__main__":