Update ui/ui_core.py

ui/ui_core.py

@@ -75,7 +75,7 @@ def convert_file_to_json(file_path: str, file_type: str) -> str:
     except Exception as e:
         return json.dumps({"error": f"Error reading {os.path.basename(file_path)}: {str(e)}"})
 
-def chunk_text(text: str, max_tokens: int = 8192) -> List[str]:
+def chunk_text(text: str, max_tokens: int = 6000) -> List[str]:
     chunks = []
     words = text.split()
     chunk = []
@@ -94,9 +94,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;'>\U0001F4CB CPS: Clinical Patient Support System</h1>")
+        gr.Markdown("<h1 style='text-align: center;'>📋 CPS: Clinical Patient Support System</h1>")
 
-        chatbot = gr.Chatbot(label="CPS Assistant", height=600, type="tuples")
+        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"],
@@ -117,18 +117,18 @@ def create_ui(agent: TxAgent):
             )
 
             try:
-                history.append((message, "⏳ Processing your request..."))
+                history.append(("user", message))
+                history.append(("assistant", "⏳ Processing your request..."))
                 yield history
 
                 extracted_text = ""
                 if uploaded_files and isinstance(uploaded_files, list):
-                    total_files = len(uploaded_files)
-                    for index, file in enumerate(uploaded_files):
+                    for file in uploaded_files:
                         if not hasattr(file, 'name'):
                             continue
                         path = file.name
-                        extension = path.split(".")[-1].lower()
-                        json_text = convert_file_to_json(path, extension)
+                        ext = path.split(".")[-1].lower()
+                        json_text = convert_file_to_json(path, ext)
                         extracted_text += sanitize_utf8(json_text) + "\n"
 
                 chunks = chunk_text(extracted_text.strip())
@@ -138,44 +138,46 @@ def create_ui(agent: TxAgent):
                         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=chunked_prompt,
-                        history=[],
-                        temperature=0.3,
-                        max_new_tokens=1024,
-                        max_token=8192,
-                        call_agent=False,
-                        conversation=conversation,
-                        uploaded_files=uploaded_files,
-                        max_round=30
-                    )
-                    result = ""
-                    for update in generator:
-                        if isinstance(update, str):
-                            result += update
-                        elif isinstance(update, list):
-                            for msg in update:
-                                if hasattr(msg, 'content'):
-                                    result += msg.content
-                    return result
-
-                # ⏱ Parallel Execution for Speed
+                    try:
+                        generator = agent.run_gradio_chat(
+                            message=chunked_prompt,
+                            history=[],
+                            temperature=0.3,
+                            max_new_tokens=1024,
+                            max_token=8192,
+                            call_agent=False,
+                            conversation=conversation,
+                            uploaded_files=uploaded_files,
+                            max_round=30
+                        )
+                        result = ""
+                        for update in generator:
+                            if update is None:
+                                print(f"[Warning] Empty response in chunk {i+1}")
+                                continue
+                            if isinstance(update, str):
+                                result += update
+                            elif isinstance(update, list):
+                                for msg in update:
+                                    if hasattr(msg, 'content'):
+                                        result += msg.content
+                        return result if result.strip() else f"[Chunk {i+1}] ⚠️ No response received."
+                    except Exception as err:
+                        print(f"[Error in chunk {i+1}] {err}")
+                        return f"[Chunk {i+1}] ❌ Failed to process due to error."
+
                 with ThreadPoolExecutor(max_workers=min(8, len(chunks))) as executor:
                     futures = [executor.submit(process_chunk, i, chunk) for i, chunk in enumerate(chunks)]
                     results = [f.result() for f in as_completed(futures)]
 
                 full_response = "\n\n".join(results)
                 full_response = clean_final_response(full_response.strip())
-                history[-1] = (message, full_response)
+                history[-1] = ("assistant", full_response)
                 yield history
 
             except Exception as chat_error:
                 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))
+                history[-1] = ("assistant", "❌ An error occurred while processing your request.")
                 yield history
 
         inputs = [message_input, chatbot, conversation_state, file_upload]