Update app.py

app.py

@@ -277,15 +277,57 @@ def init_agent():
     return agent
 
 def create_ui(agent):
-    with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    with gr.Blocks(theme=gr.themes.Soft(), title="Clinical Oversight Assistant") as demo:
         gr.Markdown("<h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>")
-        chatbot = gr.Chatbot(label="Detailed Analysis", height=600, type="messages")
-        final_summary = gr.Markdown(label="Summary of Missed Diagnoses")
-        file_upload = gr.File(file_types=[".pdf", ".csv", ".xls", ".xlsx"], file_count="multiple")
-        msg_input = gr.Textbox(placeholder="Ask about potential oversights...", show_label=False)
-        send_btn = gr.Button("Analyze", variant="primary")
-        download_output = gr.File(label="Download Full Report")
-        progress_bar = gr.Progress()
+        
+        with gr.Row():
+            with gr.Column(scale=3):
+                chatbot = gr.Chatbot(
+                    label="Analysis Conversation",
+                    height=600,
+                    bubble_full_width=False,
+                    show_copy_button=True,
+                    avatar_images=(
+                        "assets/user.png", 
+                        "assets/assistant.png"
+                    )
+                )
+            with gr.Column(scale=1):
+                final_summary = gr.Markdown(
+                    label="Summary of Findings",
+                    value="### Summary will appear here\nAfter analysis completes"
+                )
+                download_output = gr.File(
+                    label="Download Full Report",
+                    visible=False
+                )
+        
+        with gr.Row():
+            file_upload = gr.File(
+                file_types=[".pdf", ".csv", ".xls", ".xlsx"],
+                file_count="multiple",
+                label="Upload Patient Records"
+            )
+        
+        with gr.Row():
+            msg_input = gr.Textbox(
+                placeholder="Ask about potential oversights...",
+                show_label=False,
+                container=False,
+                scale=7,
+                autofocus=True
+            )
+            send_btn = gr.Button(
+                "Analyze",
+                variant="primary",
+                scale=1,
+                min_width=100
+            )
+        
+        progress_bar = gr.Progress(
+            label="Processing Progress",
+            visible=False
+        )
 
         prompt_template = """
 Analyze the patient record excerpt for missed diagnoses only. Provide a concise, evidence-based summary as a single paragraph without headings or bullet points. Include specific clinical findings (e.g., 'elevated blood pressure (160/95) on page 10'), their potential implications (e.g., 'may indicate untreated hypertension'), and a recommendation for urgent review. Do not include other oversight categories like medication conflicts. If no missed diagnoses are found, state 'No missed diagnoses identified' in a single sentence.
@@ -293,9 +335,37 @@ Patient Record Excerpt (Chunk {0} of {1}):
 {chunk}
 """
 
-        def analyze(message: str, history: List[dict], files: List, progress=gr.Progress()):
+        def process_response_stream(prompt: str, history: List[dict]) -> Generator[dict, None, None]:
+            """Process a single prompt and stream the response"""
+            full_response = ""
+            for chunk_output in agent.run_gradio_chat(prompt, [], 0.2, 512, 2048, False, []):
+                if chunk_output is None:
+                    continue
+                
+                if isinstance(chunk_output, list):
+                    for m in chunk_output:
+                        if hasattr(m, 'content') and m.content:
+                            cleaned = clean_response(m.content)
+                            if cleaned:
+                                full_response += cleaned + " "
+                                yield {"role": "assistant", "content": full_response}
+                elif isinstance(chunk_output, str) and chunk_output.strip():
+                    cleaned = clean_response(chunk_output)
+                    if cleaned:
+                        full_response += cleaned + " "
+                        yield {"role": "assistant", "content": full_response}
+            
+            return full_response
+
+        def analyze(message: str, history: List[dict], files: List) -> Generator[dict, None, None]:
+            # Start with user message
             history.append({"role": "user", "content": message})
-            yield history, None, ""
+            yield {
+                "chatbot": history,
+                "download_output": None,
+                "final_summary": "",
+                "progress_bar": gr.Progress(visible=True)
+            }
 
             extracted = []
             file_hash_value = ""
@@ -306,11 +376,7 @@ Patient Record Excerpt (Chunk {0} of {1}):
                     futures = []
                     for f in files:
                         file_type = f.name.split(".")[-1].lower()
-                        futures.append(executor.submit(
-                            process_file,
-                            f.name,
-                            file_type
-                        ))
+                        futures.append(executor.submit(process_file, f.name, file_type))
                     
                     for future in as_completed(futures):
                         try:
@@ -321,7 +387,12 @@ Patient Record Excerpt (Chunk {0} of {1}):
                 
                 file_hash_value = file_hash(files[0].name) if files else ""
                 history.append({"role": "assistant", "content": "✅ File processing complete"})
-                yield history, None, ""
+                yield {
+                    "chatbot": history,
+                    "download_output": None,
+                    "final_summary": "",
+                    "progress_bar": gr.Progress(0.2, visible=True, label="Processing files")
+                }
 
             # Convert extracted data to JSON text
             text_content = "\n".join(json.dumps(item) for item in extracted)
@@ -329,56 +400,45 @@ Patient Record Excerpt (Chunk {0} of {1}):
             # Tokenize and chunk the content properly
             chunks = tokenize_and_chunk(text_content)
             combined_response = ""
-            batch_size = 2  # Reduced batch size to prevent token overflow
             
             try:
-                for batch_idx in range(0, len(chunks), batch_size):
-                    batch_chunks = chunks[batch_idx:batch_idx + batch_size]
-                    batch_prompts = [
-                        prompt_template.format(
-                            batch_idx + i + 1,
-                            len(chunks),
-                            chunk=chunk[:1800]  # Conservative chunk size
-                        )
-                        for i, chunk in enumerate(batch_chunks)
-                    ]
+                for chunk_idx, chunk in enumerate(chunks, 1):
+                    prompt = prompt_template.format(chunk_idx, len(chunks), chunk=chunk[:1800])
                     
-                    progress((batch_idx) / len(chunks), 
-                           desc=f"Analyzing batch {(batch_idx // batch_size) + 1}/{(len(chunks) + batch_size - 1) // batch_size}")
+                    # Create a placeholder message
+                    history.append({"role": "assistant", "content": ""})
+                    yield {
+                        "chatbot": history,
+                        "download_output": None,
+                        "final_summary": "",
+                        "progress_bar": gr.Progress(
+                            0.2 + (chunk_idx/len(chunks))*0.7, 
+                            visible=True, 
+                            label=f"Analyzing chunk {chunk_idx}/{len(chunks)}"
+                        )
+                    }
                     
-                    # Process batch in parallel
-                    with ThreadPoolExecutor(max_workers=len(batch_prompts)) as executor:
-                        future_to_prompt = {
-                            executor.submit(
-                                agent.run_gradio_chat,
-                                prompt, [], 0.2, 512, 2048, False, []
-                            ): prompt
-                            for prompt in batch_prompts
+                    # Process and stream the response
+                    chunk_response = ""
+                    for update in process_response_stream(prompt, history):
+                        # Update the last message with streaming content
+                        history[-1] = update
+                        chunk_response = update["content"]
+                        yield {
+                            "chatbot": history,
+                            "download_output": None,
+                            "final_summary": "",
+                            "progress_bar": gr.Progress(
+                                0.2 + (chunk_idx/len(chunks))*0.7, 
+                                visible=True
+                            )
                         }
-                        
-                        for future in as_completed(future_to_prompt):
-                            chunk_response = ""
-                            for chunk_output in future.result():
-                                if chunk_output is None:
-                                    continue
-                                if isinstance(chunk_output, list):
-                                    for m in chunk_output:
-                                        if hasattr(m, 'content') and m.content:
-                                            cleaned = clean_response(m.content)
-                                            if cleaned:
-                                                chunk_response += cleaned + " "
-                                elif isinstance(chunk_output, str) and chunk_output.strip():
-                                    cleaned = clean_response(chunk_output)
-                                    if cleaned:
-                                        chunk_response += cleaned + " "
-                            
-                            combined_response += f"--- Analysis for Chunk {batch_idx + 1} ---\n{chunk_response.strip()}\n"
-                            history[-1] = {"role": "assistant", "content": combined_response.strip()}
-                            yield history, None, ""
-                            
-                            # Clean up memory
-                            torch.cuda.empty_cache()
-                            gc.collect()
+                    
+                    combined_response += f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response}\n"
+                    
+                    # Clean up memory
+                    torch.cuda.empty_cache()
+                    gc.collect()
 
                 # Generate final summary
                 summary = summarize_findings(combined_response)
@@ -387,15 +447,53 @@ Patient Record Excerpt (Chunk {0} of {1}):
                     with open(report_path, "w", encoding="utf-8") as f:
                         f.write(combined_response + "\n\n" + summary)
                 
-                yield history, report_path if report_path and os.path.exists(report_path) else None, summary
+                yield {
+                    "chatbot": history,
+                    "download_output": gr.File(report_path) if report_path and os.path.exists(report_path) else None,
+                    "final_summary": summary,
+                    "progress_bar": gr.Progress(1.0, visible=False)
+                }
 
             except Exception as e:
                 logger.error("Analysis error: %s", e)
                 history.append({"role": "assistant", "content": f"❌ Error occurred: {str(e)}"})
-                yield history, None, f"Error occurred during analysis: {str(e)}"
-
-        send_btn.click(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output, final_summary])
-        msg_input.submit(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output, final_summary])
+                yield {
+                    "chatbot": history,
+                    "download_output": None,
+                    "final_summary": f"Error occurred during analysis: {str(e)}",
+                    "progress_bar": gr.Progress(visible=False)
+                }
+
+        def clear_and_start():
+            return {
+                "chatbot": [],
+                "download_output": None,
+                "final_summary": "",
+                "msg_input": "",
+                "file_upload": None
+            }
+
+        # Event handlers
+        send_btn.click(
+            analyze,
+            inputs=[msg_input, chatbot, file_upload],
+            outputs=[chatbot, download_output, final_summary, progress_bar],
+            show_progress="hidden"
+        )
+        
+        msg_input.submit(
+            analyze,
+            inputs=[msg_input, chatbot, file_upload],
+            outputs=[chatbot, download_output, final_summary, progress_bar],
+            show_progress="hidden"
+        )
+        
+        demo.load(
+            clear_and_start,
+            outputs=[chatbot, download_output, final_summary, msg_input, file_upload],
+            queue=False
+        )
+    
     return demo
 
 if __name__ == "__main__":
@@ -403,13 +501,20 @@ if __name__ == "__main__":
         logger.info("Launching app...")
         agent = init_agent()
         demo = create_ui(agent)
-        demo.queue(api_open=False).launch(
+        demo.queue(
+            api_open=False,
+            max_size=20
+        ).launch(
             server_name="0.0.0.0",
             server_port=7860,
             show_error=True,
             allowed_paths=[report_dir],
-            share=False
+            share=False,
+            favicon_path="assets/favicon.ico"
         )
+    except Exception as e:
+        logger.error(f"Failed to launch app: {e}")
+        raise
     finally:
         if torch.distributed.is_initialized():
             torch.distributed.destroy_process_group()