Update ui/ui_core.py

ui/ui_core.py

@@ -12,25 +12,7 @@ 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 re.sub(r'[\ud800-\udfff]', '', text)
-
-def chunk_text(text: str, max_tokens=8000) -> List[str]:
-    chunks = []
-    lines = text.split("\n")
-    current_chunk = []
-    current_tokens = 0
-    for line in lines:
-        line_tokens = len(line.split())
-        if current_tokens + line_tokens > max_tokens:
-            chunks.append("\n".join(current_chunk))
-            current_chunk = [line]
-            current_tokens = line_tokens
-        else:
-            current_chunk.append(line)
-            current_tokens += line_tokens
-    if current_chunk:
-        chunks.append("\n".join(current_chunk))
-    return chunks
+    return text.encode("utf-8", "ignore").decode("utf-8")
 
 def extract_all_text_from_csv_or_excel(file_path: str, progress=None, index=0, total=1) -> str:
     try:
@@ -78,6 +60,23 @@ def extract_all_text_from_pdf(file_path: str, progress=None, index=0, total=1) -
     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;'>📋 CPS: Clinical Patient Support System</h1>")
@@ -109,6 +108,7 @@ def create_ui(agent: TxAgent):
                     for index, file in enumerate(uploaded_files):
                         if not hasattr(file, 'name'):
                             continue
+
                         path = file.name
                         try:
                             if path.endswith((".csv", ".xls", ".xlsx")):
@@ -122,15 +122,14 @@ def create_ui(agent: TxAgent):
                             continue
 
                 sanitized = sanitize_utf8(extracted_text.strip())
-                chunks = chunk_text(sanitized, max_tokens=8000)
+                chunks = chunk_text(sanitized, max_tokens=8192)
 
                 for i, chunk in enumerate(chunks):
-                    chunked_prompt = (
-                        f"{context}\n\n--- Uploaded File Content (Chunk {i+1}/{len(chunks)}) ---\n\n{chunk}\n\n--- End of Chunk ---\n\nNow begin your reasoning:"
+                    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:"
                     )
-
                     generator = agent.run_gradio_chat(
-                        message=chunked_prompt,
+                        message=full_message,
                         history=history,
                         temperature=0.3,
                         max_new_tokens=1024,
@@ -142,20 +141,14 @@ def create_ui(agent: TxAgent):
                     )
 
                     for update in generator:
-                        try:
-                            if isinstance(update, list):
-                                cleaned = [msg for msg in update if hasattr(msg, 'role') and hasattr(msg, 'content')]
-                                if cleaned:
-                                    yield cleaned
-                            elif isinstance(update, str):
-                                yield sanitize_utf8(update.encode("utf-8", "replace").decode("utf-8"))
-                        except Exception as update_error:
-                            print(f"Error processing update: {update_error}")
-                            continue
+                        if isinstance(update, list):
+                            yield update
+                        elif isinstance(update, str):
+                            yield [("assistant", sanitize_utf8(update))]
 
             except Exception as chat_error:
                 print(f"Chat handling error: {chat_error}")
-                yield "An error occurred while processing your request. Please try again."
+                yield [("assistant", "An error occurred while processing your request. Please try again.")]
 
         inputs = [message_input, chatbot, conversation_state, file_upload]
         send_button.click(fn=handle_chat, inputs=inputs, outputs=chatbot)