Update ui/ui_core.py

ui/ui_core.py

@@ -12,9 +12,11 @@ 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:
-    # Remove surrogate pairs and re-encode safely
-    clean_text = re.sub(r'[\ud800-\udfff]', '', text)
-    return clean_text.encode("utf-8", "ignore").decode("utf-8", "ignore")
+    return re.sub(r'[\ud800-\udfff]', '', text)
+
+def chunk_text(text: str, chunk_size: int = 8000) -> List[str]:
+    words = text.split()
+    return [" ".join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]
 
 def extract_all_text_from_csv_or_excel(file_path: str, progress=None, index=0, total=1) -> str:
     try:
@@ -36,11 +38,10 @@ 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)
-        content = f"📄 {os.path.basename(file_path)}\n\n" + "\n".join(lines)
-        return sanitize_utf8(content)
+        return f"📄 {os.path.basename(file_path)}\n\n" + "\n".join(lines)
 
     except Exception as e:
-        return sanitize_utf8(f"[Error reading {os.path.basename(file_path)}]: {str(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:
     try:
@@ -58,11 +59,10 @@ 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)}")
-        content = f"📄 {os.path.basename(file_path)}\n\n" + "\n\n".join(extracted)
-        return sanitize_utf8(content)
+        return f"📄 {os.path.basename(file_path)}\n\n" + "\n\n".join(extracted)
 
     except Exception as e:
-        return sanitize_utf8(f"[Error reading PDF {os.path.basename(file_path)}]: {str(e)}")
+        return f"[Error reading PDF {os.path.basename(file_path)}]: {str(e)}"
 
 def create_ui(agent: TxAgent):
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
@@ -103,46 +103,68 @@ def create_ui(agent: TxAgent):
                             elif path.endswith(".pdf"):
                                 extracted_text += extract_all_text_from_pdf(path, progress, index, total_files) + "\n"
                             else:
-                                extracted_text += sanitize_utf8(f"(Uploaded file: {os.path.basename(path)})\n")
+                                extracted_text += f"(Uploaded file: {os.path.basename(path)})\n"
                         except Exception as file_error:
-                            extracted_text += sanitize_utf8(f"[Error processing file: {os.path.basename(path)}] — {str(file_error)}\n")
+                            extracted_text += f"[Error processing file: {os.path.basename(path)}] — {str(file_error)}\n"
                             continue
 
-                    message = sanitize_utf8(
-                        f"{context}\n\n--- Uploaded File Content ---\n\n{extracted_text.strip()}\n\n--- End of File ---\n\nNow begin your reasoning:"
+                    sanitized_text = sanitize_utf8(extracted_text.strip())
+                    chunks = chunk_text(sanitized_text)
+
+                    for i, chunk in enumerate(chunks):
+                        message_chunk = (
+                            f"{context}\n\n--- Uploaded File Chunk {i + 1} ---\n\n{chunk}\n\n--- End of Chunk ---\n\nNow begin your reasoning:"
+                        )
+                        generator = agent.run_gradio_chat(
+                            message=message_chunk,
+                            history=history,
+                            temperature=0.3,
+                            max_new_tokens=1024,
+                            max_token=8192,
+                            call_agent=False,
+                            conversation=conversation,
+                            uploaded_files=uploaded_files,
+                            max_round=30
+                        )
+
+                        for update in generator:
+                            try:
+                                if isinstance(update, list):
+                                    cleaned = [
+                                        msg for msg in update
+                                        if hasattr(msg, 'role') and not (
+                                            msg.role == "assistant"
+                                            and hasattr(msg, 'content')
+                                            and msg.content.strip().startswith("🧠")
+                                        )
+                                    ]
+                                    if cleaned:
+                                        yield cleaned
+                                elif isinstance(update, str) and not update.strip().startswith("🧠"):
+                                    yield sanitize_utf8(update.encode("utf-8", "replace").decode("utf-8"))
+                            except Exception as update_error:
+                                print(f"Error processing update: {update_error}")
+                                continue
+
+                else:
+                    # Fallback for message-only interactions
+                    generator = agent.run_gradio_chat(
+                        message=message,
+                        history=history,
+                        temperature=0.3,
+                        max_new_tokens=1024,
+                        max_token=8192,
+                        call_agent=False,
+                        conversation=conversation,
+                        uploaded_files=uploaded_files,
+                        max_round=30
                     )
 
-                generator = agent.run_gradio_chat(
-                    message=message,
-                    history=history,
-                    temperature=0.3,
-                    max_new_tokens=1024,
-                    max_token=8192,
-                    call_agent=False,
-                    conversation=conversation,
-                    uploaded_files=uploaded_files,
-                    max_round=30
-                )
-
-                for update in generator:
-                    try:
-                        if isinstance(update, list):
-                            cleaned = [
-                                msg for msg in update
-                                if hasattr(msg, 'role')
-                                and not (
-                                    msg.role == "assistant"
-                                    and hasattr(msg, 'content')
-                                    and sanitize_utf8(msg.content.strip()).startswith("🧠")
-                                )
-                            ]
-                            if cleaned:
-                                yield cleaned
-                        elif isinstance(update, str) and not sanitize_utf8(update.strip()).startswith("🧠"):
-                            yield sanitize_utf8(update)
-                    except Exception as update_error:
-                        print(f"Error processing update: {update_error}")
-                        continue
+                    for update in generator:
+                        if isinstance(update, str):
+                            yield sanitize_utf8(update.encode("utf-8", "replace").decode("utf-8"))
+                        else:
+                            yield update
 
             except Exception as chat_error:
                 print(f"Chat handling error: {chat_error}")
@@ -158,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