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
+from typing import List, Tuple, Dict, Any, Generator, Union
 import hashlib
 import shutil
 import re
@@ -126,8 +126,13 @@ def init_agent():
     return agent
 
 
-def analyze_with_agent(agent, prompt: str) -> str:
-    try:
+def stream_final_report(agent, file) -> Generator[Union[Dict[str, str], Tuple[List[Dict[str, str]], str]], None, None]:
+    extracted_text = extract_text_from_excel(file.name)
+    chunks = split_text_into_chunks(extracted_text)
+    chunk_responses = []
+
+    for chunk in chunks:
+        prompt = build_prompt_from_text(chunk)
         response = ""
         for result in agent.run_gradio_chat(
             message=prompt,
@@ -138,45 +143,44 @@ def analyze_with_agent(agent, prompt: str) -> str:
             call_agent=False,
             conversation=[],
         ):
-            if isinstance(result, list):
+            if isinstance(result, str):
+                response += result
+            elif hasattr(result, "content"):
+                response += result.content
+            elif isinstance(result, list):
                 for r in result:
-                    if hasattr(r, 'content') and r.content:
-                        response += clean_response(r.content) + "\n"
-            elif isinstance(result, str):
-                response += clean_response(result) + "\n"
-            elif hasattr(result, 'content'):
-                response += clean_response(result.content) + "\n"
-        return response.strip()
-    except Exception as e:
-        return f"Error in analysis: {str(e)}"
-
-
-def analyze(file):
-    if not file:
-        raise gr.Error("Please upload a file")
-
-    try:
-        extracted_text = extract_text_from_excel(file.name)
-        chunks = split_text_into_chunks(extracted_text)
-
-        chunk_responses = []
-        for chunk in chunks:
-            prompt = build_prompt_from_text(chunk)
-            chunk_responses.append(analyze_with_agent(agent, prompt))
-
-        final_prompt = "\n\n".join(chunk_responses) + "\n\nSummarize the key findings above."
-        final_response = analyze_with_agent(agent, final_prompt)
-
-        full_report = f"# \U0001f9e0 Final Patient Report\n\n{final_response}"
-
-        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(full_report)
-
-        return [("user", f"[Excel Uploaded: {file.name}]"), ("assistant", full_report)], report_path
-
-    except Exception as e:
-        raise gr.Error(f"Error: {str(e)}")
+                    if hasattr(r, "content"):
+                        response += r.content
+        chunk_responses.append(clean_response(response))
+
+    final_prompt = "\n\n".join(chunk_responses) + "\n\nSummarize the key findings above."
+    yield {"role": "user", "content": f"[Excel Uploaded: {file.name}]"}
+    stream_text = ""
+    for result in agent.run_gradio_chat(
+        message=final_prompt,
+        history=[],
+        temperature=0.2,
+        max_new_tokens=MAX_NEW_TOKENS,
+        max_token=MAX_TOKENS,
+        call_agent=False,
+        conversation=[],
+    ):
+        if isinstance(result, str):
+            stream_text += result
+        elif hasattr(result, "content"):
+            stream_text += result.content
+        elif isinstance(result, list):
+            for r in result:
+                if hasattr(r, "content"):
+                    stream_text += r.content
+        yield {"role": "assistant", "content": clean_response(stream_text)}
+
+    final_report = f"# \U0001f9e0 Final Patient Report\n\n{clean_response(stream_text)}"
+    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(final_report)
+
+    yield [{"role": "user", "content": f"[Excel Uploaded: {file.name}]"}, {"role": "assistant", "content": final_report}], report_path
 
 
 def create_ui(agent):
@@ -186,8 +190,8 @@ def create_ui(agent):
         analyze_btn = gr.Button("🧠 Analyze Patient History")
         report_output = gr.File(label="Download Report")
 
-        analyze_btn.click(
-            analyze,
+        analyze_btn.stream(
+            fn=stream_final_report,
             inputs=[file_upload],
             outputs=[chatbot, report_output]
         )
@@ -207,4 +211,4 @@ if __name__ == "__main__":
         )
     except Exception as e:
         print(f"Error: {str(e)}")
-        sys.exit(1)
+        sys.exit(1)