Update ui/ui_core.py

ui/ui_core.py

@@ -3,7 +3,8 @@ import os
 import pandas as pd
 import pdfplumber
 import gradio as gr
-from typing import List
+import re
+from typing import List, Dict, Optional
 
 # ✅ Fix: Add src to Python path
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "src")))
@@ -11,12 +12,43 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..",
 from txagent.txagent import TxAgent
 
 def sanitize_utf8(text: str) -> str:
-    return text.encode("utf-8", "ignore").decode("utf-8")
-
-def clean_final_response(text: str) -> str:
-    return text.replace("[TOOL_CALLS]", "").strip()
+    """Clean text of problematic Unicode characters"""
+    return text.encode('utf-8', 'ignore').decode('utf-8')
+
+def clean_final_response(response: str) -> str:
+    """Remove tool calls and other artifacts from final response"""
+    # Split on TOOL_CALLS if present
+    if '[TOOL_CALLS]' in response:
+        response = response.split('[TOOL_CALLS]')[0]
+    # Remove any remaining special tokens
+    response = re.sub(r'\[[A-Z_]+\]', '', response)
+    return response.strip()
+
+def chunk_text(text: str, max_tokens: int = 8000) -> List[str]:
+    """Split text into chunks based on token count estimate"""
+    words = text.split()
+    chunks = []
+    current_chunk = []
+    current_tokens = 0
+    
+    for word in words:
+        # Estimate tokens (roughly 1 token per 4 characters)
+        word_tokens = len(word) // 4 + 1
+        if current_tokens + word_tokens > max_tokens and current_chunk:
+            chunks.append(' '.join(current_chunk))
+            current_chunk = [word]
+            current_tokens = word_tokens
+        else:
+            current_chunk.append(word)
+            current_tokens += word_tokens
+    
+    if current_chunk:
+        chunks.append(' '.join(current_chunk))
+    
+    return chunks
 
 def extract_all_text_from_csv_or_excel(file_path: str, progress=None, index=0, total=1) -> str:
+    """Extract text from spreadsheet files with error handling"""
     try:
         if not os.path.exists(file_path):
             return f"File not found: {file_path}"
@@ -36,12 +68,13 @@ 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"\U0001F4C4 {os.path.basename(file_path)}\n\n" + "\n".join(lines)
+        return f"📄 {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)}"
 
 def extract_all_text_from_pdf(file_path: str, progress=None, index=0, total=1) -> str:
+    """Extract text from PDF files with error handling"""
     try:
         if not os.path.exists(file_path):
             return f"PDF not found: {file_path}"
@@ -54,42 +87,31 @@ def extract_all_text_from_pdf(file_path: str, progress=None, index=0, total=1) -
                     text = page.extract_text() or ""
                     extracted.append(text.strip())
                     if progress:
-                        progress((index + (i / num_pages)) / total, desc=f"Reading PDF: {os.path.basename(file_path)} ({i+1}/{num_pages})")
+                        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"\U0001F4C4 {os.path.basename(file_path)}\n\n" + "\n\n".join(extracted)
+        return f"📄 {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)}"
 
-def chunk_text(text: str, max_tokens: int = 8192) -> List[str]:
-    chunks = []
-    words = text.split()
-    chunk = []
-    token_count = 0
-    for word in words:
-        token_count += len(word) // 4 + 1
-        if token_count > max_tokens:
-            chunks.append(" ".join(chunk))
-            chunk = [word]
-            token_count = len(word) // 4 + 1
-        else:
-            chunk.append(word)
-    if chunk:
-        chunks.append(" ".join(chunk))
-    return chunks
-
 def create_ui(agent: TxAgent):
-    with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("<h1 style='text-align: center;'>\U0001F4CB CPS: Clinical Patient Support System</h1>")
-
-        chatbot = gr.Chatbot(label="CPS Assistant", height=600, type="text")
+    with gr.Blocks(theme=gr.themes.Soft(), title="Clinical Patient Support System") as demo:
+        gr.Markdown("<h1 style='text-align: center;'>📋 CPS: Clinical Patient Support System</h1>")
+        
+        # Fix: Changed type to 'messages' to match Gradio requirements
+        chatbot = gr.Chatbot(label="CPS Assistant", height=600, type="messages")
+        
         file_upload = gr.File(
             label="Upload Medical File",
             file_types=[".pdf", ".txt", ".docx", ".jpg", ".png", ".csv", ".xls", ".xlsx"],
             file_count="multiple"
         )
-        message_input = gr.Textbox(placeholder="Ask a biomedical question or just upload the files...", show_label=False)
+        message_input = gr.Textbox(
+            placeholder="Ask a biomedical question or just upload the files...", 
+            show_label=False
+        )
         send_button = gr.Button("Send", variant="primary")
         conversation_state = gr.State([])
 
@@ -104,6 +126,7 @@ def create_ui(agent: TxAgent):
             )
 
             try:
+                # Show processing message immediately
                 history.append((message, "⏳ Processing your request..."))
                 yield history
 
@@ -146,18 +169,23 @@ def create_ui(agent: TxAgent):
                         max_round=30
                     )
 
+                    # Collect all updates from the generator
                     chunk_response = ""
                     for update in generator:
                         if isinstance(update, str):
                             chunk_response += update
                         elif isinstance(update, list):
+                            # Handle list of messages
                             for msg in update:
                                 if hasattr(msg, 'content'):
                                     chunk_response += msg.content
 
                     full_response += chunk_response + "\n\n"
 
+                # Clean up the final response
                 full_response = clean_final_response(full_response.strip())
+                
+                # Remove the processing message and add the final response
                 history[-1] = (message, full_response)
                 yield history