Update app.py

app.py

@@ -139,48 +139,35 @@ def init_agent():
     agent.init_model()
     return agent
 
-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:
+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>
+        <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:
 1. List potential missed diagnoses
 2. Flag any medication conflicts
 3. Note incomplete assessments
@@ -190,111 +177,51 @@ 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')])
-
-            # 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([])
+                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
 
         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
-        )
+        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)
 
     return demo
 
@@ -304,13 +231,10 @@ if __name__ == "__main__":
 
     print("Launching interface...")
     demo = create_ui(agent)
-    demo.queue(
-        concurrency_count=3,
-        api_open=False
-    ).launch(
+    demo.queue(api_open=False).launch(
         server_name="0.0.0.0",
         server_port=7860,
         show_error=True,
         allowed_paths=["/data/reports"],
         share=False
-    )
+    )