Update app.py

app.py

@@ -4,7 +4,7 @@ import pandas as pd
 import pdfplumber
 import json
 import gradio as gr
-from typing import List, Dict, Optional, Generator
+from typing import List, Dict, Optional, Generator, Any
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import hashlib
 import shutil
@@ -253,6 +253,10 @@ def summarize_findings(combined_response: str) -> str:
     
     return summary.strip()
 
+def update_progress(current, total, stage=""):
+    progress = f"{stage} - {current}/{total}" if stage else f"{current}/{total}"
+    return {"value": progress, "visible": True, "label": f"Progress: {progress}"}
+
 def init_agent():
     logger.info("Initializing model...")
     log_system_usage("Before Load")
@@ -276,6 +280,145 @@ def init_agent():
     logger.info("Agent Ready")
     return agent
 
+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[str, Any], None, None]:
+    # Initialize outputs
+    outputs = {
+        "chatbot": history.copy(),
+        "download_output": None,
+        "final_summary": "",
+        "progress_text": {"value": "Starting analysis...", "visible": True}
+    }
+    
+    try:
+        # Start with user message
+        history.append({"role": "user", "content": message})
+        outputs["chatbot"] = history
+        yield outputs
+        
+        extracted = []
+        file_hash_value = ""
+        
+        if files:
+            # Process files in parallel
+            with ThreadPoolExecutor(max_workers=4) as executor:
+                futures = []
+                for f in files:
+                    file_type = f.name.split(".")[-1].lower()
+                    futures.append(executor.submit(process_file, f.name, file_type))
+                
+                for i, future in enumerate(as_completed(futures), 1):
+                    try:
+                        extracted.extend(future.result())
+                        outputs["progress_text"] = update_progress(i, len(files), "Processing files")
+                        yield outputs
+                    except Exception as e:
+                        logger.error(f"File processing error: {e}")
+                        extracted.append({"error": f"Error processing file: {str(e)}"})
+            
+            file_hash_value = file_hash(files[0].name) if files else ""
+            history.append({"role": "assistant", "content": "✅ File processing complete"})
+            outputs.update({
+                "chatbot": history,
+                "progress_text": update_progress(len(files), len(files), "Files processed")
+            })
+            yield outputs
+
+        # Convert extracted data to JSON text
+        text_content = "\n".join(json.dumps(item) for item in extracted)
+        
+        # Tokenize and chunk the content properly
+        chunks = tokenize_and_chunk(text_content)
+        combined_response = ""
+        
+        for chunk_idx, chunk in enumerate(chunks, 1):
+            prompt = f"""
+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.
+
+Patient Record Excerpt (Chunk {chunk_idx} of {len(chunks)}):
+{chunk[:1800]}
+"""
+            
+            # Create a placeholder message
+            history.append({"role": "assistant", "content": ""})
+            outputs.update({
+                "chatbot": history,
+                "progress_text": update_progress(chunk_idx, len(chunks)), "Analyzing")
+            })
+            yield outputs
+            
+            # Process and stream the response
+            chunk_response = ""
+            for update in process_response_stream(prompt, history):
+                history[-1] = update
+                chunk_response = update["content"]
+                outputs.update({
+                    "chatbot": history,
+                    "progress_text": update_progress(chunk_idx, len(chunks)), "Analyzing")
+                })
+                yield outputs
+            
+            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)
+        report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt") if file_hash_value else None
+        if report_path:
+            with open(report_path, "w", encoding="utf-8") as f:
+                f.write(combined_response + "\n\n" + summary)
+        
+        outputs.update({
+            "download_output": report_path if report_path and os.path.exists(report_path) else None,
+            "final_summary": summary,
+            "progress_text": {"visible": False}
+        })
+        yield outputs
+
+    except Exception as e:
+        logger.error("Analysis error: %s", e)
+        history.append({"role": "assistant", "content": f"❌ Error occurred: {str(e)}"})
+        outputs.update({
+            "chatbot": history,
+            "final_summary": f"Error occurred during analysis: {str(e)}",
+            "progress_text": {"visible": False}
+        })
+        yield outputs
+
+def clear_and_start():
+    return [
+        [],  # chatbot
+        None,  # download_output
+        "",  # final_summary
+        "",  # msg_input
+        None,  # file_upload
+        {"visible": False}  # progress_text
+    ]
+
 def create_ui(agent):
     with gr.Blocks(theme=gr.themes.Soft(), title="Clinical Oversight Assistant") as demo:
         gr.Markdown("<h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>")
@@ -289,7 +432,7 @@ def create_ui(agent):
                     avatar_images=(
                         "assets/user.png", 
                         "assets/assistant.png"
-                    ),
+                    ) if os.path.exists("assets/user.png") else None,
                     render=False
                 )
             with gr.Column(scale=1):
@@ -330,151 +473,6 @@ def create_ui(agent):
             interactive=False
         )
 
-        def update_progress(current, total, stage=""):
-            progress = f"{stage} - {current}/{total}" if stage else f"{current}/{total}"
-            return {"value": progress, "visible": True, "label": f"Progress: {progress}"}
-
-        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.
-Patient Record Excerpt (Chunk {0} of {1}):
-{chunk}
-"""
-
-        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 {
-                "chatbot": history,
-                "download_output": None,
-                "final_summary": "",
-                "progress_text": {"value": "Starting analysis...", "visible": True}
-            }
-
-            extracted = []
-            file_hash_value = ""
-            
-            if files:
-                # Process files in parallel
-                with ThreadPoolExecutor(max_workers=4) as executor:
-                    futures = []
-                    for f in files:
-                        file_type = f.name.split(".")[-1].lower()
-                        futures.append(executor.submit(process_file, f.name, file_type))
-                    
-                    for i, future in enumerate(as_completed(futures), 1):
-                        try:
-                            extracted.extend(future.result())
-                            yield {
-                                "progress_text": update_progress(i, len(files), "Processing files")
-                            }
-                        except Exception as e:
-                            logger.error(f"File processing error: {e}")
-                            extracted.append({"error": f"Error processing file: {str(e)}"})
-                
-                file_hash_value = file_hash(files[0].name) if files else ""
-                history.append({"role": "assistant", "content": "✅ File processing complete"})
-                yield {
-                    "chatbot": history,
-                    "download_output": None,
-                    "final_summary": "",
-                    "progress_text": update_progress(len(files), len(files), "Files processed")
-                }
-
-            # Convert extracted data to JSON text
-            text_content = "\n".join(json.dumps(item) for item in extracted)
-            
-            # Tokenize and chunk the content properly
-            chunks = tokenize_and_chunk(text_content)
-            combined_response = ""
-            
-            try:
-                for chunk_idx, chunk in enumerate(chunks, 1):
-                    prompt = prompt_template.format(chunk_idx, len(chunks), chunk=chunk[:1800])
-                    
-                    # Create a placeholder message
-                    history.append({"role": "assistant", "content": ""})
-                    yield {
-                        "chatbot": history,
-                        "download_output": None,
-                        "final_summary": "",
-                        "progress_text": update_progress(chunk_idx, len(chunks), "Analyzing")
-                    }
-                    
-                    # 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_text": update_progress(chunk_idx, len(chunks), "Analyzing")
-                        }
-                    
-                    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)
-                report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt") if file_hash_value else None
-                if report_path:
-                    with open(report_path, "w", encoding="utf-8") as f:
-                        f.write(combined_response + "\n\n" + summary)
-                
-                yield {
-                    "chatbot": history,
-                    "download_output": report_path if report_path and os.path.exists(report_path) else None,
-                    "final_summary": summary,
-                    "progress_text": {"visible": False}
-                }
-
-            except Exception as e:
-                logger.error("Analysis error: %s", e)
-                history.append({"role": "assistant", "content": f"❌ Error occurred: {str(e)}"})
-                yield {
-                    "chatbot": history,
-                    "download_output": None,
-                    "final_summary": f"Error occurred during analysis: {str(e)}",
-                    "progress_text": {"visible": False}
-                }
-
-        def clear_and_start():
-            return [
-                [],  # chatbot
-                None,  # download_output
-                "",  # final_summary
-                "",  # msg_input
-                None,  # file_upload
-                {"visible": False}  # progress_text
-            ]
-
         # Event handlers
         send_btn.click(
             analyze,