Update app.py

app.py

@@ -3,7 +3,7 @@ import os
 import pandas as pd
 import json
 import gradio as gr
-from typing import List, Tuple, Union, Generator
+from typing import List, Tuple, Dict, Any, Union
 import hashlib
 import shutil
 import re
@@ -42,20 +42,17 @@ def clean_response(text: str) -> str:
 def estimate_tokens(text: str) -> int:
     return len(text) // 3.5 + 1
 
-def extract_text_from_excel(file_obj: Union[str, os.PathLike, 'file']) -> str:
+def extract_text_from_excel(file_path: str) -> str:
     all_text = []
-    try:
-        xls = pd.ExcelFile(file_obj)
-    except Exception as e:
-        raise ValueError(f"❌ Error reading Excel file: {e}")
+    xls = pd.ExcelFile(file_path)
     for sheet_name in xls.sheet_names:
         df = xls.parse(sheet_name).astype(str).fillna("")
-        rows = df.apply(lambda row: " | ".join([cell for cell in row if cell.strip()]), axis=1)
-        sheet_text = [f"[{sheet_name}] {line}" for line in rows if line.strip()]
+        rows = df.apply(lambda row: " | ".join(row), axis=1)
+        sheet_text = [f"[{sheet_name}] {line}" for line in rows]
         all_text.extend(sheet_text)
     return "\n".join(all_text)
 
-def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS, max_chunks: int = 30) -> List[str]:
+def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS) -> List[str]:
     effective_max = max_tokens - PROMPT_OVERHEAD
     lines, chunks, curr_chunk, curr_tokens = text.split("\n"), [], [], 0
     for line in lines:
@@ -63,13 +60,11 @@ def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS, max_ch
         if curr_tokens + t > effective_max:
             if curr_chunk:
                 chunks.append("\n".join(curr_chunk))
-            if len(chunks) >= max_chunks:
-                break
             curr_chunk, curr_tokens = [line], t
         else:
             curr_chunk.append(line)
             curr_tokens += t
-    if curr_chunk and len(chunks) < max_chunks:
+    if curr_chunk:
         chunks.append("\n".join(curr_chunk))
     return chunks
 
@@ -92,48 +87,14 @@ Analyze the following clinical notes and provide a detailed, concise summary foc
 Respond in well-structured bullet points with medical reasoning.
 """
 
-def validate_tool_file(file_path):
-    try:
-        with open(file_path, 'r') as f:
-            data = json.load(f)
-        if isinstance(data, list):
-            assert all(isinstance(t, dict) and "name" in t for t in data), "Invalid list format"
-        elif isinstance(data, dict):
-            assert "tools" in data and isinstance(data["tools"], list), "'tools' field missing or invalid"
-            assert all(isinstance(t, dict) and "name" in t for t in data["tools"]), "Invalid item in 'tools'"
-        else:
-            raise ValueError("Unexpected structure")
-        return True
-    except Exception as e:
-        print(f"❌ Tool validation failed for {file_path}: {e}")
-        return False
-
 def init_agent():
-    all_tool_paths = {
-        "opentarget": "/home/user/.pyenv/versions/3.10.17/lib/python3.10/site-packages/tooluniverse/data/opentarget_tools.json",
-        "fda_drug_label": "/home/user/.pyenv/versions/3.10.17/lib/python3.10/site-packages/tooluniverse/data/fda_drug_labeling_tools.json",
-        "special_tools": "/home/user/.pyenv/versions/3.10.17/lib/python3.10/site-packages/tooluniverse/data/special_tools.json",
-        "monarch": "/home/user/.pyenv/versions/3.10.17/lib/python3.10/site-packages/tooluniverse/data/monarch_tools.json",
-        "new_tool": os.path.join(tool_cache_dir, "new_tool.json"),
-    }
-
-    if not os.path.exists(all_tool_paths["new_tool"]):
-        shutil.copy(os.path.abspath("data/new_tool.json"), all_tool_paths["new_tool"])
-
-    valid_tool_paths = {}
-    for key, path in all_tool_paths.items():
-        if validate_tool_file(path):
-            valid_tool_paths[key] = path
-        else:
-            print(f"⚠️ Skipping invalid tool file: {path}")
-
-    if not valid_tool_paths:
-        raise RuntimeError("❌ No valid tool files found to load into TxAgent.")
-
+    tool_path = os.path.join(tool_cache_dir, "new_tool.json")
+    if not os.path.exists(tool_path):
+        shutil.copy(os.path.abspath("data/new_tool.json"), tool_path)
     agent = TxAgent(
         model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
         rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
-        tool_files_dict=valid_tool_paths,
+        tool_files_dict={"new_tool": tool_path},
         force_finish=True,
         enable_checker=True,
         step_rag_num=4,
@@ -142,111 +103,126 @@ def init_agent():
     agent.init_model()
     return agent
 
-def stream_report(agent, input_file: Union[str, 'file'], full_output: str) -> Generator[Tuple[str, Union[str, None], str], None, None]:
-    accumulated_text = ""
-    try:
-        if input_file is None:
-            yield "❌ Please upload a valid Excel file.", None, ""
-            return
+def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], Union[str, None]]:
+    messages = chatbot_state if chatbot_state else []
+    if file is None or not hasattr(file, "name"):
+        return messages + [{"role": "assistant", "content": "❌ Please upload a valid Excel file."}], None
 
-        if hasattr(input_file, "read"):
-            text = extract_text_from_excel(input_file)
-        elif isinstance(input_file, str) and os.path.exists(input_file):
-            text = extract_text_from_excel(input_file)
-        else:
-            raise ValueError("❌ Invalid or missing file.")
-
-        chunks = split_text_into_chunks(text)
-
-        for i, chunk in enumerate(chunks):
-            prompt = build_prompt_from_text(chunk)
-            partial = ""
-            for res in agent.run_gradio_chat(
-                message=prompt, history=[], temperature=0.2,
-                max_new_tokens=MAX_NEW_TOKENS, max_token=MAX_MODEL_TOKENS,
-                call_agent=False, conversation=[]
-            ):
-                if isinstance(res, str):
-                    partial += res
-                elif hasattr(res, "content"):
-                    partial += res.content
-            cleaned = clean_response(partial)
-            accumulated_text += f"\n\n📄 **Chunk {i+1}**:\n{cleaned}"
-            yield accumulated_text, None, ""
+    messages.append({"role": "user", "content": f"Processing Excel file: {os.path.basename(file.name)}"})
+    text = extract_text_from_excel(file.name)
+    chunks = split_text_into_chunks(text)
+    chunk_responses = [None] * len(chunks)
 
-        summary_prompt = f"Summarize this analysis in a final structured report:\n\n" + accumulated_text
-        final_report = ""
-        for res in agent.run_gradio_chat(
-            message=summary_prompt, history=[], temperature=0.2,
-            max_new_tokens=MAX_NEW_TOKENS, max_token=MAX_MODEL_TOKENS,
-            call_agent=False, conversation=[]
-        ):
+    def analyze_chunk(i, chunk):
+        prompt = build_prompt_from_text(chunk)
+        response = ""
+        for res in agent.run_gradio_chat(message=prompt, history=[], temperature=0.2, max_new_tokens=MAX_NEW_TOKENS, max_token=MAX_MODEL_TOKENS, call_agent=False, conversation=[]):
             if isinstance(res, str):
-                final_report += res
+                response += res
             elif hasattr(res, "content"):
-                final_report += res.content
-
-        cleaned = clean_response(final_report)
-        accumulated_text += f"\n\n📊 **Final Summary**:\n{cleaned}"
-        report_path = os.path.join(report_dir, f"report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md")
-        with open(report_path, 'w') as f:
-            f.write(f"# 🧠 Final Patient Report\n\n{cleaned}")
-
-        yield accumulated_text, report_path, cleaned
-
-    except Exception as e:
-        yield f"❌ Error: {str(e)}", None, ""
+                response += res.content
+            elif isinstance(res, list):
+                for r in res:
+                    if hasattr(r, "content"):
+                        response += r.content
+        return i, clean_response(response)
+
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        futures = [executor.submit(analyze_chunk, i, c) for i, c in enumerate(chunks)]
+        for f in as_completed(futures):
+            i, result = f.result()
+            chunk_responses[i] = result
+
+    valid = [r for r in chunk_responses if r and not r.startswith("❌")]
+    if not valid:
+        return messages + [{"role": "assistant", "content": "❌ No valid chunk results."}], None
+
+    summary_prompt = f"Summarize this analysis in a final structured report:\n\n" + "\n\n".join(valid)
+    messages.append({"role": "assistant", "content": "📊 Generating final report..."})
+
+    final_report = ""
+    for res in agent.run_gradio_chat(message=summary_prompt, history=[], temperature=0.2, max_new_tokens=MAX_NEW_TOKENS, max_token=MAX_MODEL_TOKENS, call_agent=False, conversation=[]):
+        if isinstance(res, str):
+            final_report += res
+        elif hasattr(res, "content"):
+            final_report += res.content
+
+    cleaned = clean_response(final_report)
+    report_path = os.path.join(report_dir, f"report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md")
+    with open(report_path, 'w') as f:
+        f.write(f"# 🧠 Final Patient Report\n\n{cleaned}")
+
+    messages.append({"role": "assistant", "content": f"📊 Final Report:\n\n{cleaned}"})
+    messages.append({"role": "assistant", "content": f"✅ Report generated and saved: {os.path.basename(report_path)}"})
+    return messages, report_path
 
 def create_ui(agent):
     with gr.Blocks(css="""
-        body {
-            background: #10141f;
-            color: #ffffff;
+        html, body, .gradio-container {
+            height: 100vh;
+            background-color: #111827;
+            color: #e5e7eb;
             font-family: 'Inter', sans-serif;
-            margin: 0;
-            padding: 0;
         }
-        .gradio-container {
-            padding: 30px;
-            width: 100vw;
-            max-width: 100%;
-            border-radius: 0;
-            background-color: #1a1f2e;
+        .message-avatar {
+            width: 38px;
+            height: 38px;
+            border-radius: 50%;
+            margin-right: 10px;
+        }
+        .chat-message {
+            display: flex;
+            align-items: flex-start;
+            margin-bottom: 1rem;
         }
-        .output-markdown {
-            background-color: #131720;
+        .message-bubble {
+            background-color: #1f2937;
+            padding: 12px 16px;
             border-radius: 12px;
-            padding: 20px;
-            min-height: 600px;
-            overflow-y: auto;
-            border: 1px solid #2c3344;
+            max-width: 90%;
         }
-        .gr-button {
-            background: linear-gradient(135deg, #4b4ced, #37b6e9);
+        .chat-input {
+            background-color: #1f2937;
+            border: 1px solid #374151;
+            border-radius: 8px;
+            color: #e5e7eb;
+            padding: 0.75rem 1rem;
+        }
+        .gr-button.primary {
+            background: #2563eb;
             color: white;
-            font-weight: 500;
-            border: none;
-            padding: 10px 20px;
             border-radius: 8px;
-            transition: background 0.3s ease;
+            font-weight: 600;
         }
-        .gr-button:hover {
-            background: linear-gradient(135deg, #37b6e9, #4b4ced);
+        .gr-button.primary:hover {
+            background: #1e40af;
         }
     """) as demo:
-        gr.Markdown("""# 🧠 Clinical Reasoning Assistant  
-Upload clinical Excel records below and click **Analyze** to generate a medical summary.""")
-        file_upload = gr.File(label="Upload Excel File", file_types=[".xlsx"])
-        analyze_btn = gr.Button("Analyze")
-        report_output_markdown = gr.Markdown(elem_classes="output-markdown")
-        report_file = gr.File(label="Download Report", visible=False)
-        full_output = gr.State(value="")
-
-        analyze_btn.click(
-            fn=stream_report,
-            inputs=[file_upload, full_output],
-            outputs=[report_output_markdown, report_file, full_output]
-        )
+        gr.Markdown("""<h2 style='color:#60a5fa'>🩺 Patient History AI Assistant</h2><p>Upload a clinical Excel file and receive a structured diagnostic summary.</p>""")
+        with gr.Row():
+            with gr.Column(scale=3):
+                chatbot = gr.Chatbot(
+                    label="Clinical Assistant",
+                    height=700,
+                    type="messages",
+                    avatar_images=[
+                        "https://ui-avatars.com/api/?name=AI&background=2563eb&color=fff&size=128",
+                        "https://ui-avatars.com/api/?name=You&background=374151&color=fff&size=128"
+                    ]
+                )
+            with gr.Column(scale=1):
+                with gr.Row():
+                    file_upload = gr.File(label="", file_types=[".xlsx"], elem_id="upload-btn")
+                    analyze_btn = gr.Button("🧠 Analyze", variant="primary")
+                report_output = gr.File(label="Download Report", visible=False, interactive=False)
+
+        chatbot_state = gr.State(value=[])
+
+        def update_ui(file, current_state):
+            messages, report_path = process_final_report(agent, file, current_state)
+            return messages, gr.update(visible=report_path is not None, value=report_path), messages
+
+        analyze_btn.click(fn=update_ui, inputs=[file_upload, chatbot_state], outputs=[chatbot, report_output, chatbot_state])
 
     return demo
 
@@ -254,7 +230,7 @@ if __name__ == "__main__":
     try:
         agent = init_agent()
         demo = create_ui(agent)
-        demo.launch(server_name="0.0.0.0", server_port=7860, allowed_paths=["/data/hf_cache/reports"], share=True)
+        demo.launch(server_name="0.0.0.0", server_port=7860, allowed_paths=["/data/hf_cache/reports"], share=False)
     except Exception as e:
         print(f"Error: {str(e)}")
         sys.exit(1)