Update ui/ui_core.py

ui/ui_core.py

@@ -3,7 +3,6 @@ import os
 import pandas as pd
 import pdfplumber
 import gradio as gr
-import re
 from typing import List
 
 # ✅ Fix: Add src to Python path
@@ -14,8 +13,8 @@ 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 clean_final_response(text: str) -> str:
+    return text.replace("[TOOL_CALLS]", "").strip()
 
 def extract_all_text_from_csv_or_excel(file_path: str, progress=None, index=0, total=1) -> str:
     try:
@@ -37,7 +36,7 @@ def extract_all_text_from_csv_or_excel(file_path: str, progress=None, index=0, t
             line = " | ".join(str(cell) for cell in row if pd.notna(cell))
             if line:
                 lines.append(line)
-        return f"📄 {os.path.basename(file_path)}\n\n" + "\n".join(lines)
+        return f"\U0001F4C4 {os.path.basename(file_path)}\n\n" + "\n".join(lines)
 
     except Exception as e:
         return f"[Error reading {os.path.basename(file_path)}]: {str(e)}"
@@ -58,7 +57,7 @@ def extract_all_text_from_pdf(file_path: str, progress=None, index=0, total=1) -
                         progress((index + (i / num_pages)) / total, desc=f"Reading PDF: {os.path.basename(file_path)} ({i+1}/{num_pages})")
                 except Exception as e:
                     extracted.append(f"[Error reading page {i+1}]: {str(e)}")
-        return f"📄 {os.path.basename(file_path)}\n\n" + "\n\n".join(extracted)
+        return f"\U0001F4C4 {os.path.basename(file_path)}\n\n" + "\n\n".join(extracted)
 
     except Exception as e:
         return f"[Error reading PDF {os.path.basename(file_path)}]: {str(e)}"
@@ -82,9 +81,9 @@ def chunk_text(text: str, max_tokens: int = 8192) -> List[str]:
 
 def create_ui(agent: TxAgent):
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("<h1 style='text-align: center;'>📋 CPS: Clinical Patient Support System</h1>")
+        gr.Markdown("<h1 style='text-align: center;'>\U0001F4CB CPS: Clinical Patient Support System</h1>")
 
-        chatbot = gr.Chatbot(label="CPS Assistant", height=600, type="messages", show_copy_button=True)
+        chatbot = gr.Chatbot(label="CPS Assistant", height=600, type="text")
         file_upload = gr.File(
             label="Upload Medical File",
             file_types=[".pdf", ".txt", ".docx", ".jpg", ".png", ".csv", ".xls", ".xlsx"],
@@ -105,7 +104,8 @@ def create_ui(agent: TxAgent):
             )
 
             try:
-                yield history + [{"role": "assistant", "content": "⏳ **Processing... Please wait while I analyze the files.**"}]
+                history.append((message, "⏳ Processing your request..."))
+                yield history
 
                 extracted_text = ""
                 if uploaded_files and isinstance(uploaded_files, list):
@@ -127,13 +127,15 @@ def create_ui(agent: TxAgent):
                 sanitized = sanitize_utf8(extracted_text.strip())
                 chunks = chunk_text(sanitized)
 
-                all_responses = ""
+                full_response = ""
                 for i, chunk in enumerate(chunks):
-                    full_message = (
-                        f"{context}\n\n--- Uploaded File Chunk {i+1}/{len(chunks)} ---\n\n{chunk}\n\n--- End of Chunk ---\n\nNow begin your reasoning:"
+                    chunked_prompt = (
+                        f"{context}\n\n--- Uploaded File Content (Chunk {i+1}/{len(chunks)}) ---\n\n{chunk}\n\n"
+                        f"--- End of Chunk ---\n\nNow begin your analysis:"
                     )
+
                     generator = agent.run_gradio_chat(
-                        message=full_message,
+                        message=chunked_prompt,
                         history=[],
                         temperature=0.3,
                         max_new_tokens=1024,
@@ -143,24 +145,30 @@ def create_ui(agent: TxAgent):
                         uploaded_files=uploaded_files,
                         max_round=30
                     )
+
+                    chunk_response = ""
                     for update in generator:
                         if isinstance(update, str):
-                            all_responses += update + "\n"
+                            chunk_response += update
+                        elif isinstance(update, list):
+                            for msg in update:
+                                if hasattr(msg, 'content'):
+                                    chunk_response += msg.content
+
+                    full_response += chunk_response + "\n\n"
 
-                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
+                full_response = clean_final_response(full_response.strip())
+                history[-1] = (message, full_response)
+                yield history
 
             except Exception as chat_error:
-                print(f"Chat error: {chat_error}")
-                final_history = history + [
-                    {"role": "user", "content": message},
-                    {"role": "assistant", "content": "❌ An error occurred while processing your request."}
-                ]
-                yield final_history
+                print(f"Chat handling error: {chat_error}")
+                error_msg = "An error occurred while processing your request. Please try again."
+                if len(history) > 0 and history[-1][1].startswith("⏳"):
+                    history[-1] = (history[-1][0], error_msg)
+                else:
+                    history.append((message, error_msg))
+                yield history
 
         inputs = [message_input, chatbot, conversation_state, file_upload]
         send_button.click(fn=handle_chat, inputs=inputs, outputs=chatbot)