Update ui/ui_core.py

ui/ui_core.py

@@ -4,17 +4,20 @@ import sys
 import pandas as pd
 import pdfplumber
 
-# Add src to Python path
+# ✅ Add src to Python path
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "src")))
 from txagent.txagent import TxAgent
 
 
-def extract_all_text_from_csv(file_path):
+def extract_all_text_from_csv_or_excel(file_path):
     try:
-        df = pd.read_csv(file_path, low_memory=False)
+        if file_path.endswith(".xlsx"):
+            df = pd.read_excel(file_path, engine="openpyxl")
+        else:
+            df = pd.read_csv(file_path, low_memory=False)
         return df.to_string(index=False)
     except Exception as e:
-        return f"Error parsing CSV: {e}"
+        return f"Error parsing file: {e}"
 
 
 def extract_all_text_from_pdf(file_path):
@@ -36,25 +39,30 @@ def create_ui(agent: TxAgent):
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown("<h1 style='text-align: center;'>🧠 CPS: Clinical Processing System</h1>")
         chatbot = gr.Chatbot(label="CPS Assistant", height=600, type="messages")
+        conversation_state = gr.State([])
 
-        # Hidden file upload, attached to input bar
-        with gr.Row():
+        with gr.Row(equal_height=True):
             uploaded_files = gr.File(
-                label="📎", file_types=[".pdf", ".txt", ".docx", ".jpg", ".png", ".csv"],
-                file_count="multiple", visible=False
+                visible=False,
+                file_types=[".pdf", ".txt", ".docx", ".jpg", ".png", ".csv", ".xlsx"],
+                file_count="multiple"
             )
 
             with gr.Column(scale=10):
                 message_input = gr.Textbox(
-                    placeholder="Type your medical question or upload files...", show_label=False, scale=10
+                    placeholder="Type your medical question or upload files...",
+                    show_label=False,
+                    container=True,
+                    scale=10
                 )
 
-            with gr.Column(scale=1, min_width=60):
-                file_icon = gr.UploadButton("📎", file_types=[".pdf", ".csv", ".docx", ".txt", ".jpg", ".png"], file_count="multiple")
-
-            send_button = gr.Button("Send", variant="primary")
+            with gr.Column(scale=1, min_width=50):
+                file_icon = gr.UploadButton(
+                    "📎", file_types=[".pdf", ".txt", ".docx", ".csv", ".xlsx", ".jpg", ".png"],
+                    file_count="multiple"
+                )
 
-        conversation_state = gr.State([])
+            send_button = gr.Button("Send", variant="primary", scale=1)
 
         def handle_chat(message, history, conversation, new_files):
             context = (
@@ -67,8 +75,8 @@ def create_ui(agent: TxAgent):
                 extracted_text = ""
                 for file in new_files:
                     path = file.name
-                    if path.endswith(".csv"):
-                        extracted_text += extract_all_text_from_csv(path) + "\n"
+                    if path.endswith((".csv", ".xlsx")):
+                        extracted_text += extract_all_text_from_csv_or_excel(path) + "\n"
                     elif path.endswith(".pdf"):
                         extracted_text += extract_all_text_from_pdf(path) + "\n"
                     else:
@@ -90,11 +98,14 @@ def create_ui(agent: TxAgent):
             for update in generator:
                 yield update
 
-        # Bind send logic
         file_icon.upload(fn=None, inputs=[], outputs=[uploaded_files])
         send_button.click(fn=handle_chat, inputs=[message_input, chatbot, conversation_state, uploaded_files], outputs=chatbot)
         message_input.submit(fn=handle_chat, inputs=[message_input, chatbot, conversation_state, uploaded_files], outputs=chatbot)
 
-        gr.Examples([["Upload your medical form and ask what the doctor might’ve missed."]], inputs=message_input)
+        gr.Examples([
+            ["Upload your medical form and ask what the doctor might’ve missed."],
+            ["This patient was treated with antibiotics for UTI. What else should we check?"],
+            ["Is there anything abnormal in the attached blood work report?"]
+        ], inputs=message_input)
 
     return demo