Update app.py

app.py

@@ -130,7 +130,7 @@ def init_agent():
         enable_checker=True,
         step_rag_num=8,
         seed=100,
-        additional_default_tools=[],
+        additional_default_tools=[]
     )
     agent.init_model()
     return agent
@@ -140,35 +140,28 @@ def create_ui(agent: TxAgent):
         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)
-        file_upload = gr.File(
-            label="Upload Medical Records",
-            file_types=[".pdf", ".csv", ".xls", ".xlsx"],
-            file_count="multiple"
-        )
+        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:
-                # Initialize conversation
-                history.append((message, "Analyzing records for potential oversights..."))
+                history.append({"role": "user", "content": message})
+                history.append({"role": "assistant", "content": "Analyzing records for potential oversights..."})
                 yield history, None
 
-                # Process files
                 extracted_data = ""
                 file_hash_value = ""
-                if files and isinstance(files, list):
+                if files:
                     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 ""
+                        futures = [executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower()) for f in files if hasattr(f, 'name')]
+                        results = [sanitize_utf8(f.result()) for f in as_completed(futures)]
+                        extracted_data = "\n".join(results)
+                        file_hash_value = file_hash(files[0].name)
 
-                # Medical oversight analysis prompt
                 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
@@ -177,16 +170,10 @@ def create_ui(agent: TxAgent):
 
 Medical Records:\n{extracted_data[:15000]}
 
-Provide ONLY the potential oversights in this format:
-
-### Potential Oversights:
-1. [Missed diagnosis] - [Evidence from records]
-2. [Medication issue] - [Supporting data]
-3. [Assessment gap] - [Relevant findings]"""
+### Potential Oversights:\n"""
 
-                # Generate and stream response
-                full_response = ""
-                generator = agent.run_gradio_chat(
+                response = []
+                for chunk in agent.run_gradio_chat(
                     message=analysis_prompt,
                     history=[],
                     temperature=0.2,
@@ -194,45 +181,26 @@ Provide ONLY the potential oversights in this format:
                     max_token=4096,
                     call_agent=False,
                     conversation=conversation
-                )
-
-                for update in generator:
-                    if not update:
-                        continue
-                    
-                    if isinstance(update, str):
-                        full_response += update
-                    elif isinstance(update, list):
-                        full_response += "".join([msg.content for msg in update if hasattr(msg, 'content')])
-                    
-                    # Clean and update the response
-                    cleaned = full_response.replace("[TOOL_CALLS]", "").strip()
-                    if cleaned:
-                        history[-1] = (message, cleaned)
-                        yield history, None
-
-                # Final cleaned response
-                final_output = full_response.replace("[TOOL_CALLS]", "").strip()
+                ):
+                    if isinstance(chunk, str):
+                        response.append(chunk)
+                    elif isinstance(chunk, list):
+                        response.extend([c.content for c in chunk if hasattr(c, 'content')])
+                    history[-1] = {"role": "assistant", "content": "".join(response).strip()}
+                    yield history, None
+
+                final_output = "".join(response).strip()
                 if not final_output:
                     final_output = "No clear oversights identified. Recommend comprehensive review."
+                history[-1] = {"role": "assistant", "content": final_output}
 
-                # Prepare report path if available
-                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[-1] = (message, final_output)
-                print(f"Final analysis:\n{final_output}")
-                yield history, report_path
+                report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt")
+                return history, report_path if os.path.exists(report_path) else None
 
             except Exception as e:
-                print(f"Analysis error: {str(e)}")
-                history[-1] = (message, f"❌ Analysis failed: {str(e)}")
-                yield history, None
+                history.append({"role": "assistant", "content": f"❌ Analysis failed: {str(e)}"})
+                return history, None
 
-        # UI event handlers
         inputs = [msg_input, chatbot, conversation_state, file_upload]
         outputs = [chatbot, download_output]
         send_btn.click(analyze_potential_oversights, inputs=inputs, outputs=outputs)
@@ -249,7 +217,7 @@ Provide ONLY the potential oversights in this format:
 if __name__ == "__main__":
     print("Initializing medical analysis agent...")
     agent = init_agent()
-    
+
     print("Performing warm-up call...")
     try:
         warm_up = agent.run_gradio_chat(
@@ -268,7 +236,7 @@ if __name__ == "__main__":
 
     print("Launching interface...")
     demo = create_ui(agent)
-    demo.queue(concurrency_count=2).launch(
+    demo.queue().launch(
         server_name="0.0.0.0",
         server_port=7860,
         show_error=True,