Update ui/ui_core.py

ui/ui_core.py

@@ -3,6 +3,7 @@ import os
 import pandas as pd
 import pdfplumber
 import gradio as gr
+import re
 from typing import List
 
 # ✅ Fix: Add src to Python path
@@ -13,6 +14,9 @@ from txagent.txagent import TxAgent
 def sanitize_utf8(text: str) -> str:
     return text.encode("utf-8", "ignore").decode("utf-8")
 
+def clean_final_response(response: str) -> str:
+    return re.split(r"\\[TOOL_CALLS\\]", response)[0].strip()
+
 def extract_all_text_from_csv_or_excel(file_path: str, progress=None, index=0, total=1) -> str:
     try:
         if not os.path.exists(file_path):
@@ -95,18 +99,14 @@ def create_ui(agent: TxAgent):
                 "You are an expert clinical AI assistant reviewing medical form or interview data. "
                 "Your job is to analyze this data and reason about any information or red flags that a human doctor might have overlooked. "
                 "Provide a **detailed and structured response**, including examples, supporting evidence from the form, and clinical rationale for why these items matter. "
+                "Ensure the output is informative and helpful for improving patient care. "
                 "Do not hallucinate. Base the response only on the provided form content. "
                 "End with a section labeled '🧠 Final Analysis' where you summarize key findings the doctor may have missed."
             )
 
-            # Show centered loading message in chatbot
-            updated_history = history + [
-                {"role": "user", "content": message},
-                {"role": "assistant", "content": "<div style='text-align:center'>⏳ Processing... Please wait while I analyze the files.</div>"}
-            ]
-            yield updated_history
-
             try:
+                yield history + [{"role": "assistant", "content": "⏳ **Processing... Please wait while I analyze the files.**"}]
+
                 extracted_text = ""
                 if uploaded_files and isinstance(uploaded_files, list):
                     total_files = len(uploaded_files)
@@ -145,21 +145,22 @@ def create_ui(agent: TxAgent):
                     )
                     for update in generator:
                         if isinstance(update, str):
-                            all_responses += update
-
-                all_responses = sanitize_utf8(all_responses.strip())
+                            all_responses += update + "\n"
 
-                # Replace the temporary loading message with the final answer
-                updated_history[-1] = {"role": "assistant", "content": all_responses}
-                yield updated_history
+                all_responses = clean_final_response(sanitize_utf8(all_responses.strip()))
+                final_history = history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": all_responses}
+                ]
+                yield final_history
 
             except Exception as chat_error:
                 print(f"Chat error: {chat_error}")
-                updated_history[-1] = {
-                    "role": "assistant",
-                    "content": "❌ An error occurred while processing your request. Please try again."
-                }
-                yield updated_history
+                final_history = history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": "❌ An error occurred while processing your request."}
+                ]
+                yield final_history
 
         inputs = [message_input, chatbot, conversation_state, file_upload]
         send_button.click(fn=handle_chat, inputs=inputs, outputs=chatbot)
@@ -171,4 +172,4 @@ def create_ui(agent: TxAgent):
             ["Is there anything abnormal in the attached blood work report?"]
         ], inputs=message_input)
 
-    return demo
+    return demo