Update app.py

app.py

@@ -139,33 +139,48 @@ def init_agent():
     agent.init_model()
     return agent
 
-def create_ui(agent: TxAgent):
-    with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("<h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>")
-        gr.Markdown("<h3 style='text-align: center;'>Identify potential oversights in patient care</h3>")
-
-        chatbot = gr.Chatbot(label="Analysis", height=600, type="messages")
-        file_upload = gr.File(label="Upload Medical Records", 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")
-        conversation_state = gr.State([])
-        download_output = gr.File(label="Download Full Report")
-
-        def analyze_potential_oversights(message: str, history: list, conversation: list, files: list):
-            start_time = time.time()
-            try:
-                history = history + [{"role": "user", "content": message}, {"role": "assistant", "content": "⏳ Analyzing records for potential oversights..."}]
-                yield history, None
-
-                extracted_data = ""
-                file_hash_value = ""
-                if files and isinstance(files, list):
-                    with ThreadPoolExecutor(max_workers=4) as executor:
-                        futures = [executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower()) for f in files if hasattr(f, 'name')]
-                        extracted_data = "\n".join([sanitize_utf8(f.result()) for f in as_completed(futures)])
-                        file_hash_value = file_hash(files[0].name) if files else ""
-
-                analysis_prompt = f"""Review these medical records and identify EXACTLY what might have been missed:
+def format_analysis_output(text: str) -> str:
+    """Clean and format the analysis output for better display"""
+    # Remove tool call artifacts
+    text = text.replace("[TOOL_CALLS]", "").strip()
+    if "[TOOL_CALLS]" in text:
+        text = text.split("[TOOL_CALLS]")[0].strip()
+    
+    # Remove duplicate sections
+    if "Based on the medical records provided" in text:
+        parts = text.split("Based on the medical records provided")
+        if len(parts) > 1:
+            text = "Based on the medical records provided" + parts[-1]
+    
+    # Format sections with Markdown and emojis
+    replacements = {
+        "1. **Missed Diagnoses**:": "### 🔍 Missed Diagnoses",
+        "2. **Medication Conflicts**:": "\n### 💊 Medication Conflicts",
+        "3. **Incomplete Assessments**:": "\n### 📋 Incomplete Assessments",
+        "4. **Abnormal Results Needing Follow-up**:": "\n### ⚠️ Abnormal Results Needing Follow-up",
+        "Overall, the patient's medical records": "\n### 📝 Overall Assessment"
+    }
+    
+    for old, new in replacements.items():
+        text = text.replace(old, new)
+    
+    return text
+
+def analyze_potential_oversights(message: str, history: list, conversation: list, files: list):
+    start_time = time.time()
+    try:
+        history = history + [{"role": "user", "content": message}, {"role": "assistant", "content": "⏳ Analyzing records for potential oversights..."}]
+        yield history, None
+
+        extracted_data = ""
+        file_hash_value = ""
+        if files and isinstance(files, list):
+            with ThreadPoolExecutor(max_workers=4) as executor:
+                futures = [executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower()) for f in files if hasattr(f, 'name')]
+                extracted_data = "\n".join([sanitize_utf8(f.result()) for f in as_completed(futures)])
+                file_hash_value = file_hash(files[0].name) if files else ""
+
+        analysis_prompt = f"""Review these medical records and identify EXACTLY what might have been missed:
 1. List potential missed diagnoses
 2. Flag any medication conflicts
 3. Note incomplete assessments
@@ -175,51 +190,111 @@ Medical Records:\n{extracted_data[:15000]}
 
 ### Potential Oversights:\n"""
 
-                response = ""
-                for chunk in agent.run_gradio_chat(
-                    message=analysis_prompt,
-                    history=[],
-                    temperature=0.2,
-                    max_new_tokens=1024,
-                    max_token=4096,
-                    call_agent=False,
-                    conversation=conversation
-                ):
-                    if isinstance(chunk, str):
-                        response += chunk
-                    elif isinstance(chunk, list):
-                        response += "".join([c.content for c in chunk if hasattr(c, 'content')])
-
-                    cleaned = response.replace("[TOOL_CALLS]", "").strip()
-                    yield history[:-1] + [{"role": "assistant", "content": cleaned}], None
-
-                final_output = response.replace("[TOOL_CALLS]", "").strip()
-                if not final_output:
-                    final_output = "No clear oversights identified. Recommend comprehensive review."
-
-                report_path = None
-                if file_hash_value:
-                    possible_report = os.path.join(report_dir, f"{file_hash_value}_report.txt")
-                    if os.path.exists(possible_report):
-                        report_path = possible_report
-
-                history = history[:-1] + [{"role": "assistant", "content": final_output}]
-                yield history, report_path
-
-            except Exception as e:
-                history.append({"role": "assistant", "content": f"❌ Analysis failed: {str(e)}"})
-                yield history, None
+        response = ""
+        for chunk in agent.run_gradio_chat(
+            message=analysis_prompt,
+            history=[],
+            temperature=0.2,
+            max_new_tokens=1024,
+            max_token=4096,
+            call_agent=False,
+            conversation=conversation
+        ):
+            if isinstance(chunk, str):
+                response += chunk
+            elif isinstance(chunk, list):
+                response += "".join([c.content for c in chunk if hasattr(c, 'content')])
+
+            # Format the partial response for display
+            formatted = format_analysis_output(response)
+            yield history[:-1] + [{"role": "assistant", "content": formatted}], None
+
+        # Final formatting
+        final_output = format_analysis_output(response)
+        if not final_output:
+            final_output = "No clear oversights identified. Recommend comprehensive review."
+
+        report_path = None
+        if file_hash_value:
+            possible_report = os.path.join(report_dir, f"{file_hash_value}_report.txt")
+            if os.path.exists(possible_report):
+                report_path = possible_report
+
+        history = history[:-1] + [{"role": "assistant", "content": final_output}]
+        yield history, report_path
+
+    except Exception as e:
+        history.append({"role": "assistant", "content": f"❌ Analysis failed: {str(e)}"})
+        yield history, None
+
+def create_ui(agent: TxAgent):
+    with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 900px !important}") as demo:
+        gr.Markdown("""
+        <div style='text-align: center;'>
+            <h1>🩺 Clinical Oversight Assistant</h1>
+            <h3>Identify potential oversights in patient care</h3>
+            <p>Upload medical records to analyze for missed diagnoses, medication conflicts, incomplete assessments, and abnormal results needing follow-up.</p>
+        </div>
+        """)
+
+        with gr.Row():
+            with gr.Column(scale=2):
+                file_upload = gr.File(
+                    label="Upload Medical Records", 
+                    file_types=[".pdf", ".csv", ".xls", ".xlsx"], 
+                    file_count="multiple",
+                    height=100
+                )
+                msg_input = gr.Textbox(
+                    placeholder="Ask about potential oversights...", 
+                    show_label=False,
+                    lines=3,
+                    max_lines=6
+                )
+                send_btn = gr.Button("Analyze", variant="primary", size="lg")
+                
+                gr.Examples(
+                    examples=[
+                        ["What might have been missed in this patient's treatment?"],
+                        ["Are there any medication conflicts in these records?"],
+                        ["What abnormal results require follow-up?"],
+                        ["Identify any incomplete assessments in these records"]
+                    ],
+                    inputs=msg_input,
+                    label="Example Queries"
+                )
+
+            with gr.Column(scale=3):
+                chatbot = gr.Chatbot(
+                    label="Analysis Results", 
+                    height=600,
+                    bubble_full_width=False,
+                    show_copy_button=True,
+                    avatar_images=(
+                        "assets/user.png", 
+                        "assets/doctor.png"
+                    )
+                )
+                download_output = gr.File(
+                    label="Download Full Report",
+                    visible=False
+                )
+
+        conversation_state = gr.State([])
 
         inputs = [msg_input, chatbot, conversation_state, file_upload]
         outputs = [chatbot, download_output]
-        send_btn.click(analyze_potential_oversights, inputs=inputs, outputs=outputs)
-        msg_input.submit(analyze_potential_oversights, inputs=inputs, outputs=outputs)
-
-        gr.Examples([
-            ["What might have been missed in this patient's treatment?"],
-            ["Are there any medication conflicts in these records?"],
-            ["What abnormal results require follow-up?"]
-        ], inputs=msg_input)
+        
+        send_btn.click(
+            analyze_potential_oversights, 
+            inputs=inputs, 
+            outputs=outputs
+        )
+        msg_input.submit(
+            analyze_potential_oversights, 
+            inputs=inputs, 
+            outputs=outputs
+        )
 
     return demo
 
@@ -229,9 +304,13 @@ if __name__ == "__main__":
 
     print("Launching interface...")
     demo = create_ui(agent)
-    demo.queue().launch(
+    demo.queue(
+        concurrency_count=3,
+        api_open=False
+    ).launch(
         server_name="0.0.0.0",
         server_port=7860,
         show_error=True,
-        allowed_paths=["/data/reports"]
-    )
+        allowed_paths=["/data/reports"],
+        share=False
+    )