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, Dict, Any, Generator, Union
+from typing import List, Tuple, Dict, Any, Union
 import hashlib
 import shutil
 import re
@@ -119,21 +119,17 @@ def init_agent():
     agent.init_model()
     return agent
 
-def stream_final_report(agent, file) -> Generator[Tuple[List[Dict[str, str]], Dict[str, Any]], None, None]:
-    messages = []
+def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], Dict[str, Any]]:
+    messages = chatbot_state if chatbot_state else []
     report_output = {"visible": False, "value": None}
 
     if file is None or not hasattr(file, "name"):
-        messages = [{"role": "assistant", "content": "❌ Please upload a valid Excel file before analyzing."}]
-        yield messages, report_output
-        return
+        messages.append({"role": "assistant", "content": "❌ Please upload a valid Excel file before analyzing."})
+        return messages, report_output
 
     try:
-        messages = [
-            {"role": "user", "content": f"Processing Excel file: {os.path.basename(file.name)}"},
-            {"role": "assistant", "content": "⏳ Extracting and analyzing data..."}
-        ]
-        yield messages, report_output
+        messages.append({"role": "user", "content": f"Processing Excel file: {os.path.basename(file.name)}"})
+        messages.append({"role": "assistant", "content": "⏳ Extracting and analyzing data..."})
 
         extracted_text = extract_text_from_excel(file.name)
         chunks = split_text_into_chunks(extracted_text)
@@ -141,7 +137,6 @@ def stream_final_report(agent, file) -> Generator[Tuple[List[Dict[str, str]], Di
 
         for i, chunk in enumerate(chunks):
             messages.append({"role": "assistant", "content": f"🔍 Analyzing chunk {i+1}/{len(chunks)}..."})
-            yield messages, report_output
             
             prompt = build_prompt_from_text(chunk)
             response = ""
@@ -165,11 +160,9 @@ def stream_final_report(agent, file) -> Generator[Tuple[List[Dict[str, str]], Di
             
             chunk_responses.append(clean_response(response))
             messages.append({"role": "assistant", "content": f"✅ Chunk {i+1} analysis complete"})
-            yield messages, report_output
 
         final_prompt = "\n\n".join(chunk_responses) + "\n\nSummarize the key findings above."
         messages.append({"role": "assistant", "content": "📊 Generating final report..."})
-        yield messages, report_output
 
         stream_text = ""
         for result in agent.run_gradio_chat(
@@ -189,11 +182,10 @@ def stream_final_report(agent, file) -> Generator[Tuple[List[Dict[str, str]], Di
                 for r in result:
                     if hasattr(r, "content"):
                         stream_text += r.content
-            
-            messages[-1]["content"] = f"📊 Generating final report...\n\n{clean_response(stream_text)}"
-            yield messages, report_output
-
+        
         final_report = f"# \U0001f9e0 Final Patient Report\n\n{clean_response(stream_text)}"
+        messages[-1]["content"] = f"📊 Final Report:\n\n{clean_response(stream_text)}"
+
         timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
         report_path = os.path.join(report_dir, f"report_{timestamp}.md")
         
@@ -202,11 +194,11 @@ def stream_final_report(agent, file) -> Generator[Tuple[List[Dict[str, str]], Di
 
         messages.append({"role": "assistant", "content": f"✅ Report generated and saved: report_{timestamp}.md"})
         report_output = {"visible": True, "value": report_path}
-        yield messages, report_output
 
     except Exception as e:
         messages.append({"role": "assistant", "content": f"❌ Error processing file: {str(e)}"})
-        yield messages, report_output
+
+    return messages, report_output
 
 def create_ui(agent):
     with gr.Blocks(title="Patient History Chat", css=".gradio-container {max-width: 900px !important}") as demo:
@@ -240,10 +232,17 @@ def create_ui(agent):
                     interactive=False
                 )
 
+        # State to maintain chatbot messages
+        chatbot_state = gr.State(value=[])
+
+        def update_ui(file, current_state):
+            messages, report_output = process_final_report(agent, file, current_state)
+            return messages, report_output, messages
+
         analyze_btn.click(
-            fn=lambda file: stream_final_report(agent, file),
-            inputs=[file_upload],
-            outputs=[chatbot, report_output],
+            fn=update_ui,
+            inputs=[file_upload, chatbot_state],
+            outputs=[chatbot, report_output, chatbot_state],
             api_name="analyze"
         )