Update app.py

app.py

@@ -126,20 +126,63 @@ def init_agent():
     return agent
 
 
-def stream_final_report(agent, file) -> Generator[Tuple[List[Dict[str, str]], str], None, None]:
+def stream_final_report(agent, file) -> Generator[Tuple[List[Dict[str, str]], Union[str, None]], None, None]:
+    # Initialize with empty values
+    messages = []
+    report_path = None
+    
     if file is None or not hasattr(file, "name"):
-        yield ([{"role": "assistant", "content": "❌ Please upload a valid Excel file before analyzing."}], "")
+        messages = [{"role": "assistant", "content": "❌ Please upload a valid Excel file before analyzing."}]
+        yield messages, None
         return
 
-    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 = ""
+    try:
+        # Initial processing message
+        messages = [{"role": "user", "content": f"Processing Excel file: {file.name}"},
+                   {"role": "assistant", "content": "⏳ Extracting and analyzing data..."}]
+        yield messages, None
+
+        extracted_text = extract_text_from_excel(file.name)
+        chunks = split_text_into_chunks(extracted_text)
+        chunk_responses = []
+
+        # Process each chunk
+        for i, chunk in enumerate(chunks):
+            messages.append({"role": "assistant", "content": f"🔍 Analyzing chunk {i+1}/{len(chunks)}..."})
+            yield messages, None
+            
+            prompt = build_prompt_from_text(chunk)
+            response = ""
+            for result in agent.run_gradio_chat(
+                message=prompt,
+                history=[],
+                temperature=0.2,
+                max_new_tokens=MAX_NEW_TOKENS,
+                max_token=MAX_TOKENS,
+                call_agent=False,
+                conversation=[],
+            ):
+                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"):
+                            response += r.content
+            
+            chunk_responses.append(clean_response(response))
+            messages.append({"role": "assistant", "content": f"✅ Chunk {i+1} analysis complete"})
+            yield messages, None
+
+        # Final summarization
+        final_prompt = "\n\n".join(chunk_responses) + "\n\nSummarize the key findings above."
+        messages.append({"role": "assistant", "content": "📊 Generating final report..."})
+        yield messages, None
+
+        stream_text = ""
         for result in agent.run_gradio_chat(
-            message=prompt,
+            message=final_prompt,
             history=[],
             temperature=0.2,
             max_new_tokens=MAX_NEW_TOKENS,
@@ -148,57 +191,84 @@ def stream_final_report(agent, file) -> Generator[Tuple[List[Dict[str, str]], st
             conversation=[],
         ):
             if isinstance(result, str):
-                response += result
+                stream_text += result
             elif hasattr(result, "content"):
-                response += result.content
+                stream_text += result.content
             elif isinstance(result, list):
                 for r in result:
                     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."
-    messages = [{"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
-        messages.append({"role": "assistant", "content": clean_response(stream_text)})
-        yield (messages, None)
-
-    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)
-
-    messages.append({"role": "assistant", "content": final_report})
-    yield (messages, report_path)
+                        stream_text += r.content
+            
+            messages[-1]["content"] = f"📊 Generating final report...\n\n{clean_response(stream_text)}"
+            yield messages, None
+
+        # Save final report
+        final_report = f"# \U0001f9e0 Final Patient 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")
+        
+        with open(report_path, 'w') as f:
+            f.write(final_report)
+
+        messages.append({"role": "assistant", "content": f"✅ Report generated and saved: report_{timestamp}.md"})
+        yield messages, report_path
+
+    except Exception as e:
+        messages.append({"role": "assistant", "content": f"❌ Error: {str(e)}"})
+        yield messages, None
+
 
 def create_ui(agent):
-    with gr.Blocks(title="Patient History Chat") as demo:
-        chatbot = gr.Chatbot(label="Clinical Assistant", show_copy_button=True, type="messages")
-        file_upload = gr.File(label="Upload Excel File", file_types=[".xlsx"])
-        analyze_btn = gr.Button("🧠 Analyze Patient History")
-        report_output = gr.File(label="Download Report")
+    with gr.Blocks(title="Patient History Chat", css=".gradio-container {max-width: 900px !important}") as demo:
+        gr.Markdown("## 🏥 Patient History Analysis Tool")
+        
+        with gr.Row():
+            with gr.Column(scale=3):
+                chatbot = gr.Chatbot(
+                    label="Clinical Assistant",
+                    show_copy_button=True,
+                    height=600,
+                    avatar_images=(
+                        None,  # User avatar
+                        "https://i.imgur.com/6wX7Zb4.png"  # Bot avatar
+                    )
+                )
+            with gr.Column(scale=1):
+                file_upload = gr.File(
+                    label="Upload Excel File",
+                    file_types=[".xlsx"],
+                    height=100
+                )
+                analyze_btn = gr.Button(
+                    "🧠 Analyze Patient History",
+                    variant="primary"
+                )
+                report_output = gr.File(
+                    label="Download Report",
+                    visible=False,
+                    interactive=False
+                )
+                gr.Examples(
+                    examples=["sample_data/sample_patient_history.xlsx"],
+                    inputs=file_upload,
+                    label="Sample Files"
+                )
 
         analyze_btn.click(
             fn=lambda file: stream_final_report(agent, file),
             inputs=[file_upload],
-            outputs=[chatbot, report_output]
+            outputs=[chatbot, report_output],
+            api_name="analyze"
+        )
+
+        def show_report(report_path):
+            if report_path:
+                return gr.File(visible=True, value=report_path)
+            return gr.File(visible=False)
+
+        demo.load(
+            lambda: None,
+            outputs=report_output
         )
 
     return demo
@@ -212,7 +282,8 @@ if __name__ == "__main__":
             server_name="0.0.0.0",
             server_port=7860,
             show_error=True,
-            allowed_paths=["/data/hf_cache/reports"]
+            allowed_paths=["/data/hf_cache/reports"],
+            share=False
         )
     except Exception as e:
         print(f"Error: {str(e)}")