Update ui/ui_core.py

ui/ui_core.py

@@ -3,8 +3,7 @@ import os
 import pandas as pd
 import pdfplumber
 import gradio as gr
-import re
-from typing import List, Dict, Optional
+from typing import List
 
 # ✅ Fix: Add src to Python path
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "src")))
@@ -12,43 +11,12 @@ 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:
-    """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
+    return text.encode("utf-8", "ignore").decode("utf-8")
+
+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:
-    """Extract text from spreadsheet files with error handling"""
     try:
         if not os.path.exists(file_path):
             return f"File not found: {file_path}"
@@ -68,13 +36,12 @@ 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)}"
 
 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}"
@@ -87,31 +54,42 @@ 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"📄 {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)}"
 
+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(), 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")
-        
+    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="tuples")
         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([])
 
@@ -126,7 +104,6 @@ def create_ui(agent: TxAgent):
             )
 
             try:
-                # Show processing message immediately
                 history.append((message, "⏳ Processing your request..."))
                 yield history
 
@@ -169,23 +146,18 @@ 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
 
@@ -208,4 +180,4 @@ def create_ui(agent: TxAgent):
             ["Is there anything abnormal in the attached blood work report?"]
         ], inputs=message_input)
 
-    return demo
+    return demo