app.py way too heavy before

app.py

@@ -1,605 +1,3 @@
-# import os
-# import re
-# import shutil
-# import traceback
-# import gradio as gr
-# from pathlib import Path
-# from histopath.agent import A1
-# from dotenv import load_dotenv
-
-# # Load environment variables
-# load_dotenv()
-
-# # Get passcode from environment
-# PASSCODE = os.getenv("GRADIO_PASSWORD")
-
-# # Initialize agent (will be created after passcode validation)
-# agent = None
-
-
-# def check_for_output_files():
-#     """Check for all files in the output directory and return their paths."""
-#     output_dir = Path("./output")
-#     if not output_dir.exists():
-#         return [], []
-    
-#     image_extensions = {".png", ".jpg", ".jpeg", ".svg", ".tif", ".tiff"}
-#     data_extensions = {".csv", ".txt", ".json", ".npy"}
-    
-#     images = []
-#     data_files = []
-    
-#     for file in output_dir.iterdir():
-#         if file.is_file():
-#             if file.suffix.lower() in image_extensions:
-#                 images.append(str(file))
-#             elif file.suffix.lower() in data_extensions:
-#                 data_files.append(str(file))
-    
-#     return images, data_files
-
-
-# def preview_uploaded_file(uploaded_file):
-#     """Preview the uploaded file - show image or file info."""
-#     if uploaded_file is None:
-#         return None, None, "No file uploaded"
-    
-#     file_path = Path(uploaded_file.name)
-#     file_ext = file_path.suffix.lower()
-    
-#     image_extensions = {".png", ".jpg", ".jpeg", ".svg", ".tif", ".tiff", ".svs"}
-    
-#     if file_ext in image_extensions:
-#         # Show image preview
-#         return uploaded_file.name, None, f"📷 Previewing: {file_path.name}"
-#     else:
-#         # Show file info
-#         file_size = Path(uploaded_file.name).stat().st_size / 1024  # KB
-#         return None, uploaded_file.name, f"📄 File: {file_path.name} ({file_size:.1f} KB)"
-
-
-# def parse_agent_output(output):
-#     """Parse agent output to extract code blocks, observations, and regular text."""
-#     # Strip out the message divider bars
-#     output = re.sub(r'={30,}\s*(Human|Ai)\s+Message\s*={30,}', '', output)
-#     output = output.strip()
-    
-#     parsed = {
-#         "type": "text",
-#         "content": output,
-#         "code": None,
-#         "observation": None,
-#         "thinking": None
-#     }
-    
-#     # Check for code execution block
-#     execute_match = re.search(r'<execute>(.*?)</execute>', output, re.DOTALL)
-#     if execute_match:
-#         parsed["type"] = "code"
-#         parsed["code"] = execute_match.group(1).strip()
-#         # Extract text before the code block (thinking/explanation)
-#         text_before = output[:execute_match.start()].strip()
-#         # Remove any think tags but keep the content
-#         text_before = re.sub(r'<think>(.*?)</think>', r'\1', text_before, flags=re.DOTALL)
-#         text_before = re.sub(r'={30,}.*?={30,}', '', text_before).strip()
-#         parsed["thinking"] = text_before if text_before else None
-#         return parsed
-    
-#     # Check for observation block
-#     observation_match = re.search(r'<observation>(.*?)</observation>', output, re.DOTALL)
-#     if observation_match:
-#         parsed["type"] = "observation"
-#         parsed["observation"] = observation_match.group(1).strip()
-#         # Extract text before observation if any
-#         text_before = output[:observation_match.start()].strip()
-#         text_before = re.sub(r'<think>(.*?)</think>', r'\1', text_before, flags=re.DOTALL)
-#         text_before = re.sub(r'={30,}.*?={30,}', '', text_before).strip()
-#         parsed["thinking"] = text_before if text_before else None
-#         return parsed
-    
-#     # Check for solution block
-#     solution_match = re.search(r'<solution>(.*?)</solution>', output, re.DOTALL)
-#     if solution_match:
-#         parsed["type"] = "solution"
-#         parsed["content"] = solution_match.group(1).strip()
-#         # Get thinking before solution
-#         text_before = output[:solution_match.start()].strip()
-#         text_before = re.sub(r'<think>(.*?)</think>', r'\1', text_before, flags=re.DOTALL)
-#         text_before = re.sub(r'={30,}.*?={30,}', '', text_before).strip()
-#         parsed["thinking"] = text_before if text_before else None
-#         return parsed
-    
-#     # Clean up any remaining tags for display
-#     cleaned = re.sub(r'<think>(.*?)</think>', r'\1', output, flags=re.DOTALL)
-#     cleaned = re.sub(r'={30,}.*?={30,}', '', cleaned).strip()
-#     parsed["content"] = cleaned
-    
-#     return parsed
-
-
-# def format_message_for_display(parsed_output):
-#     """Format parsed output into a readable message for the chatbot."""
-#     msg_parts = []
-    
-#     # Add thinking/explanation text first if present
-#     if parsed_output.get("thinking"):
-#         msg_parts.append(parsed_output["thinking"])
-    
-#     if parsed_output["type"] == "code":
-#         # Add separator if there was thinking text
-#         if parsed_output.get("thinking"):
-#             msg_parts.append("\n---\n")
-        
-#         msg_parts.append("### 💻 Executing Code\n")
-#         msg_parts.append(f"```python\n{parsed_output['code']}\n```")
-        
-#     elif parsed_output["type"] == "observation":
-#         # Add separator if there was thinking text
-#         if parsed_output.get("thinking"):
-#             msg_parts.append("\n---\n")
-        
-#         msg_parts.append("### 📊 Observation\n")
-#         msg_parts.append(f"```\n{parsed_output['observation']}\n```")
-        
-#     elif parsed_output["type"] == "solution":
-#         # Add separator if there was thinking text
-#         if parsed_output.get("thinking"):
-#             msg_parts.append("\n---\n")
-        
-#         msg_parts.append("### ✅ Solution\n")
-#         msg_parts.append(parsed_output['content'])
-        
-#     else:
-#         # For regular text, just add the content if thinking wasn't already set
-#         if not parsed_output.get("thinking"):
-#             msg_parts.append(parsed_output["content"])
-    
-#     return "\n\n".join(msg_parts)
-
-
-# def process_agent_response(prompt, uploaded_file, chatbot_history):
-#     """Process the agent response and update chatbot."""
-#     global agent
-    
-#     if agent is None:
-#         chatbot_history.append({
-#             "role": "assistant",
-#             "content": "⚠️ Please enter the passcode first to initialize the agent."
-#         })
-#         yield chatbot_history, None, None, None, None, "⚠️ Agent not initialized"
-#         return
-    
-#     if not prompt.strip() and uploaded_file is None:
-#         chatbot_history.append({
-#             "role": "assistant",
-#             "content": "⚠️ Please provide a prompt or upload a file."
-#         })
-#         yield chatbot_history, None, None, None, None, "⚠️ No input provided"
-#         return
-    
-#     # Handle file upload
-#     file_path = None
-#     file_info = ""
-#     if uploaded_file is not None:
-#         try:
-#             # Create data directory if it doesn't exist
-#             data_dir = Path("./data")
-#             data_dir.mkdir(exist_ok=True)
-            
-#             # Copy uploaded file to data directory
-#             file_name = Path(uploaded_file.name).name
-#             file_path = data_dir / file_name
-#             shutil.copy(uploaded_file.name, file_path)
-            
-#             file_info = f"\n\n📎 **Uploaded file:** `{file_path}`\n"
-            
-#             # Augment prompt with file path
-#             if prompt.strip():
-#                 prompt = f"{prompt}\n\nUploaded file path: {file_path}"
-#             else:
-#                 prompt = f"I have uploaded a file at: {file_path}. Please analyze it."
-                
-#         except Exception as e:
-#             error_msg = f"❌ Error handling file upload: {str(e)}"
-#             chatbot_history.append({
-#                 "role": "assistant",
-#                 "content": error_msg
-#             })
-#             yield chatbot_history, None, None, None, None, error_msg
-#             return
-    
-#     # Add user message to chat
-#     user_message = prompt if not file_info else f"{prompt}{file_info}"
-#     chatbot_history.append({"role": "user", "content": user_message})
-#     yield chatbot_history, None, None, None, None, "🔄 Processing..."
-    
-#     try:
-#         # Stream agent responses
-#         step_count = 0
-#         for step in agent.go_stream(prompt):
-#             step_count += 1
-#             output = step.get("output", "")
-            
-#             if output:
-#                 # Parse the output
-#                 parsed = parse_agent_output(output)
-                
-#                 # Add thinking text as separate message if present
-#                 if parsed.get("thinking"):
-#                     chatbot_history.append({
-#                         "role": "assistant",
-#                         "content": parsed["thinking"]
-#                     })
-                
-#                 # Add the block (code/observation/solution) as separate message if present
-#                 if parsed["type"] == "code" and parsed["code"]:
-#                     chatbot_history.append({
-#                         "role": "assistant",
-#                         "content": f"### 💻 Executing Code\n\n```python\n{parsed['code']}\n```"
-#                     })
-#                 elif parsed["type"] == "observation" and parsed["observation"]:
-#                     chatbot_history.append({
-#                         "role": "assistant",
-#                         "content": f"### 📊 Observation\n\n```\n{parsed['observation']}\n```"
-#                     })
-#                 elif parsed["type"] == "solution":
-#                     chatbot_history.append({
-#                         "role": "assistant",
-#                         "content": f"### ✅ Solution\n\n{parsed['content']}"
-#                     })
-#                 elif parsed["type"] == "text" and parsed["content"]:
-#                     # Only add if we haven't already added it as thinking
-#                     if not parsed.get("thinking"):
-#                         chatbot_history.append({
-#                             "role": "assistant",
-#                             "content": parsed["content"]
-#                         })
-                
-#                 # Check for output files after each step
-#                 images, data_files = check_for_output_files()
-                
-#                 # Create status message
-#                 status = f"🔄 Step {step_count}"
-#                 if parsed["type"] == "code":
-#                     status += " - Executing code..."
-#                 elif parsed["type"] == "observation":
-#                     status += " - Processing results..."
-#                 elif parsed["type"] == "solution":
-#                     status += " - Finalizing solution..."
-                
-#                 yield (
-#                     chatbot_history,
-#                     images if images else None,
-#                     data_files if data_files else None,
-#                     None,
-#                     None,
-#                     status
-#                 )
-        
-#         # Final check for files
-#         final_images, final_data = check_for_output_files()
-        
-#         # Create download links message if files were generated
-#         if final_images or final_data:
-#             download_msg = "\n\n---\n\n### 📁 Generated Files Ready for Download\n\n"
-            
-#             if final_images:
-#                 download_msg += f"**🖼️ Images ({len(final_images)})** - Available in the **Images** tab →\n"
-#                 for img_path in final_images:
-#                     img_name = Path(img_path).name
-#                     download_msg += f"- `{img_name}`\n"
-#                 download_msg += "\n"
-            
-#             if final_data:
-#                 download_msg += f"**📄 Data Files ({len(final_data)})** - Available in the **Data** tab →\n"
-#                 for data_path in final_data:
-#                     data_name = Path(data_path).name
-#                     download_msg += f"- `{data_name}`\n"
-            
-#             download_msg += "\n*Click the download button on each file in the respective tabs above.*"
-            
-#             # Add download message as separate bubble
-#             chatbot_history.append({
-#                 "role": "assistant",
-#                 "content": download_msg
-#             })
-        
-#         status = "✅ Complete"
-#         if final_images:
-#             status += f" | {len(final_images)} image(s)"
-#         if final_data:
-#             status += f" | {len(final_data)} data file(s)"
-            
-#         yield chatbot_history, final_images if final_images else None, final_data if final_data else None, None, None, status
-        
-#     except Exception as e:
-#         error_msg = f"❌ Error: {str(e)}\n\n```\n{traceback.format_exc()}\n```"
-#         chatbot_history.append({
-#             "role": "assistant",
-#             "content": error_msg
-#         })
-#         yield chatbot_history, None, None, None, None, "❌ Error occurred"
-
-
-# def validate_passcode(passcode):
-#     """Validate the passcode and initialize the agent."""
-#     global agent
-    
-#     if passcode == PASSCODE:
-#         # Initialize agent
-#         try:
-#             agent = A1()
-#             return (
-#                 gr.update(visible=False),  # Hide passcode section
-#                 gr.update(visible=True),   # Show main interface
-#                 "✅ Access granted! Agent initialized and ready."
-#             )
-#         except Exception as e:
-#             error_trace = traceback.format_exc()
-#             return (
-#                 gr.update(visible=True),
-#                 gr.update(visible=False),
-#                 f"❌ Error initializing agent:\n{str(e)}\n\n{error_trace}"
-#             )
-#     else:
-#         return (
-#             gr.update(visible=True),
-#             gr.update(visible=False),
-#             "❌ Invalid passcode. Please try again."
-#         )
-
-
-# def clear_chat():
-#     """Clear the chat history and output files."""
-#     # Clean up output directory
-#     output_dir = Path("./output")
-#     if output_dir.exists():
-#         shutil.rmtree(output_dir)
-#     output_dir.mkdir(exist_ok=True)
-    
-#     # Clean up data directory
-#     data_dir = Path("./data")
-#     if data_dir.exists():
-#         for file in data_dir.iterdir():
-#             if file.is_file():
-#                 file.unlink()
-    
-#     return [], None, None, None, None, "🗑️ Chat cleared"
-
-
-# # Create Gradio interface with custom theme
-# custom_theme = gr.themes.Soft(
-#     primary_hue="blue",
-#     secondary_hue="slate",
-#     spacing_size="sm",
-#     radius_size="md",
-# ).set(
-#     button_primary_background_fill="*primary_500",
-#     button_primary_background_fill_hover="*primary_600",
-#     block_label_text_weight="600",
-#     block_title_text_weight="600",
-# )
-
-# with gr.Blocks(title="HistoPath Agent", theme=custom_theme, css="""
-#     .gradio-container {
-#         max-width: 100% !important;
-#     }
-#     .main-header {
-#         text-align: center;
-#         padding: 1.5rem 0;
-#         background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-#         color: white;
-#         border-radius: 8px;
-#         margin-bottom: 1.5rem;
-#     }
-#     .main-header h1 {
-#         margin: 0;
-#         font-size: 2.2rem;
-#         font-weight: 700;
-#     }
-#     .main-header p {
-#         margin: 0.5rem 0 0 0;
-#         opacity: 0.95;
-#         font-size: 1.1rem;
-#     }
-#     .file-upload-box .wrap {
-#         min-width: 0 !important;
-#     }
-#     .file-upload-box .file-name {
-#         word-break: break-word !important;
-#         white-space: normal !important;
-#         overflow-wrap: break-word !important;
-#     }
-#     .tab-nav {
-#         margin-bottom: 0.5rem;
-#     }
-#     /* Better styling for code and observation blocks */
-#     .message.bot pre {
-#         background-color: #f6f8fa !important;
-#         border: 1px solid #d0d7de !important;
-#         border-radius: 6px !important;
-#         padding: 12px !important;
-#         margin: 8px 0 !important;
-#     }
-#     .message.bot h3 {
-#         margin-top: 12px !important;
-#         margin-bottom: 8px !important;
-#         font-weight: 600 !important;
-#     }
-#     .message.bot hr {
-#         border: none !important;
-#         border-top: 2px solid #e1e4e8 !important;
-#         margin: 16px 0 !important;
-#     }
-# """) as demo:
-    
-#     # Header
-#     gr.HTML("""
-#         <div class="main-header">
-#             <h1>🔬 HistoPath Agent</h1>
-#             <p>AI-Powered Histopathology Analysis Assistant</p>
-#         </div>
-#     """)
-    
-#     # Passcode section
-#     with gr.Group(visible=True) as passcode_section:
-#         gr.Markdown("### 🔐 Authentication Required")
-        
-#         with gr.Row():
-#             passcode_input = gr.Textbox(
-#                 label="Passcode",
-#                 type="password",
-#                 placeholder="Enter your passcode...",
-#                 scale=3
-#             )
-#             passcode_btn = gr.Button("🔓 Unlock", variant="primary", scale=1, size="lg")
-        
-#         passcode_status = gr.Textbox(
-#             label="Status",
-#             interactive=False,
-#             lines=2
-#         )
-    
-#     # Main interface (hidden initially)
-#     with gr.Group(visible=False) as main_interface:
-#         with gr.Row(equal_height=True):
-#             # Left column - Chat interface
-#             with gr.Column(scale=3):
-#                 chatbot = gr.Chatbot(
-#                     label="💬 Conversation",
-#                     height=550,
-#                     # type="messages",
-#                     show_label=True,
-#                     render_markdown=True,
-#                 )
-                
-#                 # Input area
-#                 with gr.Row():
-#                     with gr.Column(scale=7):
-#                         prompt_input = gr.Textbox(
-#                             label="Your Query",
-#                             placeholder="E.g., 'Caption the uploaded whole slide image' or 'Segment cells using instanseg model'",
-#                             lines=2,
-#                             max_lines=5,
-#                             show_label=False,
-#                         )
-#                     with gr.Column(scale=3):
-#                         file_upload = gr.File(
-#                             label="📎 Upload File",
-#                             file_types=[".svs", ".png", ".jpg", ".jpeg", ".tif", ".tiff", ".csv", ".txt", ".json", ".npy"],
-#                             height=75,
-#                             elem_classes="file-upload-box",
-#                         )
-                
-#                 with gr.Row():
-#                     submit_btn = gr.Button("🚀 Submit", variant="primary", scale=3, size="lg")
-#                     clear_btn = gr.Button("🗑️ Clear", scale=1, size="lg", variant="secondary")
-                
-#                 status_text = gr.Textbox(
-#                     label="Status",
-#                     interactive=False,
-#                     value="Ready",
-#                     show_label=False,
-#                     container=False,
-#                 )
-            
-#             # Right column - Outputs
-#             with gr.Column(scale=2):
-#                 with gr.Tabs():
-#                     with gr.Tab("📥 Input"):
-#                         with gr.Column():
-#                             input_image_preview = gr.Image(
-#                                 label="Input Image",
-#                                 height=400,
-#                                 show_label=False,
-#                                 container=True,
-#                             )
-#                             input_file_preview = gr.File(
-#                                 label="Input File",
-#                                 interactive=False,
-#                                 height=100,
-#                                 show_label=False,
-#                                 container=True,
-#                             )
-#                             input_status = gr.Textbox(
-#                                 value="Upload a file to preview",
-#                                 show_label=False,
-#                                 interactive=False,
-#                                 container=False,
-#                             )
-                    
-#                     with gr.Tab("🖼️ Images"):
-#                         output_gallery = gr.Gallery(
-#                             label="Generated Visualizations",
-#                             columns=1,
-#                             height=600,
-#                             object_fit="contain",
-#                             show_label=False,
-#                             show_download_button=True,
-#                         )
-                    
-#                     with gr.Tab("📄 Data"):
-#                         data_files = gr.File(
-#                             label="Generated Data Files",
-#                             file_count="multiple",
-#                             interactive=False,
-#                             height=600,
-#                             show_label=False,
-#                         )
-    
-#     # Event handlers
-#     passcode_btn.click(
-#         fn=validate_passcode,
-#         inputs=[passcode_input],
-#         outputs=[passcode_section, main_interface, passcode_status]
-#     )
-    
-#     # File upload preview
-#     file_upload.change(
-#         fn=preview_uploaded_file,
-#         inputs=[file_upload],
-#         outputs=[input_image_preview, input_file_preview, input_status]
-#     )
-    
-#     submit_btn.click(
-#         fn=process_agent_response,
-#         inputs=[prompt_input, file_upload, chatbot],
-#         outputs=[chatbot, output_gallery, data_files, input_image_preview, input_file_preview, status_text]
-#     )
-    
-#     clear_btn.click(
-#         fn=clear_chat,
-#         outputs=[chatbot, output_gallery, data_files, input_image_preview, input_file_preview, status_text]
-#     )
-    
-#     # Allow enter key to submit
-#     prompt_input.submit(
-#         fn=process_agent_response,
-#         inputs=[prompt_input, file_upload, chatbot],
-#         outputs=[chatbot, output_gallery, data_files, input_image_preview, input_file_preview, status_text]
-#     )
-
-
-# if __name__ == "__main__":
-#     # Create necessary directories
-#     Path("./data").mkdir(exist_ok=True)
-#     Path("./output").mkdir(exist_ok=True)
-    
-#     print("=" * 60)
-#     print("🔬 HistoPath Agent - Gradio Interface")
-#     print("=" * 60)
-#     # print(f"Passcode: {PASSCODE}")
-#     print("Starting server...")
-#     print("=" * 60)
-    
-#     # Launch the app
-#     #demo.launch(
-#         # server_name="0.0.0.0",
-#         # server_port=7860,  # Let Gradio auto-pick an available port
-#         # share=True,
-#         # show_error=True,
-#     #)
-#     demo.launch()
 import os
 import re
 import shutil
@@ -760,7 +158,7 @@ def format_message_for_display(parsed_output):
 
 
 def process_agent_response(prompt, uploaded_file, chatbot_history):
-    """Process the agent response and update chatbot."""
+    """Process the agent response and update chatbot - AGGRESSIVE FIX: Minimal yields."""
     global agent
     
     if agent is None:
@@ -768,7 +166,7 @@ def process_agent_response(prompt, uploaded_file, chatbot_history):
             "role": "assistant",
             "content": "⚠️ Please enter the passcode first to initialize the agent."
         })
-        yield chatbot_history, gr.update(value=None), gr.update(value=None), None, None, "⚠️ Agent not initialized"
+        yield chatbot_history, None, None, None, None, "⚠️ Agent not initialized"
         return
     
     if not prompt.strip() and uploaded_file is None:
@@ -776,7 +174,7 @@ def process_agent_response(prompt, uploaded_file, chatbot_history):
             "role": "assistant",
             "content": "⚠️ Please provide a prompt or upload a file."
         })
-        yield chatbot_history, gr.update(value=None), gr.update(value=None), None, None, "⚠️ No input provided"
+        yield chatbot_history, None, None, None, None, "⚠️ No input provided"
         return
     
     # Handle file upload
@@ -807,49 +205,46 @@ def process_agent_response(prompt, uploaded_file, chatbot_history):
                 "role": "assistant",
                 "content": error_msg
             })
-            yield chatbot_history, gr.update(value=None), gr.update(value=None), None, None, error_msg
+            yield chatbot_history, None, None, None, None, error_msg
             return
     
     # Add user message to chat
     user_message = prompt if not file_info else f"{prompt}{file_info}"
     chatbot_history.append({"role": "user", "content": user_message})
-    yield chatbot_history, gr.update(value=None), gr.update(value=None), None, None, "🔄 Processing..."
+    
+    # CRITICAL FIX: Only yield once at the start to show user message
+    yield chatbot_history, None, None, None, None, "🔄 Processing..."
     
     try:
-        # Stream agent responses
+        # CRITICAL FIX: Collect ALL steps without yielding
         step_count = 0
+        collected_outputs = []
+        
         for step in agent.go_stream(prompt):
             step_count += 1
             output = step.get("output", "")
             
             if output:
-                # Parse the output
-                parsed = parse_agent_output(output)
-                
-                # Format for display
-                formatted_message = format_message_for_display(parsed)
-                
-                # Update or append to chatbot history
-                if chatbot_history and chatbot_history[-1]["role"] == "assistant":
-                    # Update the last assistant message
-                    chatbot_history[-1]["content"] = formatted_message
-                else:
-                    # Add new assistant message
-                    chatbot_history.append({
-                        "role": "assistant",
-                        "content": formatted_message
-                    })
-                
-                # CRITICAL FIX: Don't update files during streaming
-                # Only yield chatbot updates, keep files unchanged
-                yield chatbot_history, gr.update(), gr.update(), None, None, f"🔄 Step {step_count}..."
+                collected_outputs.append(output)
         
-        # CRITICAL FIX: Only check and return files AFTER streaming completes
-        images, data = check_for_output_files()
+        # CRITICAL FIX: Process ALL collected outputs at once
+        for output in collected_outputs:
+            parsed = parse_agent_output(output)
+            formatted_message = format_message_for_display(parsed)
+            
+            # Update or append to chatbot history
+            if chatbot_history and chatbot_history[-1]["role"] == "assistant":
+                # Update the last assistant message
+                chatbot_history[-1]["content"] = formatted_message
+            else:
+                # Add new assistant message
+                chatbot_history.append({
+                    "role": "assistant",
+                    "content": formatted_message
+                })
         
-        # Use gr.update with proper value to avoid Content-Length issues
-        final_images = gr.update(value=images if images else None)
-        final_data = gr.update(value=data if data else None)
+        # CRITICAL FIX: Check files only ONCE after all processing
+        images, data = check_for_output_files()
         
         status_msg = f"✅ Complete ({step_count} steps)"
         if images:
@@ -857,7 +252,8 @@ def process_agent_response(prompt, uploaded_file, chatbot_history):
         if data:
             status_msg += f" | {len(data)} data file(s)"
         
-        yield chatbot_history, final_images, final_data, None, None, status_msg
+        # CRITICAL FIX: Final single yield with all results
+        yield chatbot_history, images, data, None, None, status_msg
         
     except Exception as e:
         error_trace = traceback.format_exc()
@@ -868,12 +264,12 @@ def process_agent_response(prompt, uploaded_file, chatbot_history):
             "content": error_msg
         })
         
-        yield chatbot_history, gr.update(value=None), gr.update(value=None), None, None, f"❌ Error: {str(e)}"
+        yield chatbot_history, None, None, None, None, f"❌ Error: {str(e)}"
 
 
 def clear_chat():
     """Clear the chat history and outputs."""
-    return [], gr.update(value=None), gr.update(value=None), None, None, "Ready"
+    return [], None, None, None, None, "Ready"
 
 
 def validate_passcode(input_passcode):
@@ -1011,7 +407,6 @@ with gr.Blocks(title="HistoPath Agent", theme=custom_theme, css="""
                 chatbot = gr.Chatbot(
                     label="💬 Conversation",
                     height=550,
-                    # type="messages",
                     show_label=True,
                     render_markdown=True,
                 )
@@ -1131,7 +526,6 @@ if __name__ == "__main__":
     print("=" * 60)
     print("🔬 HistoPath Agent - Gradio Interface")
     print("=" * 60)
-    # print(f"Passcode: {PASSCODE}")
     print("Starting server...")
     print("=" * 60)
     

app_old.py

@@ -0,0 +1,602 @@
+# import os
+# import re
+# import shutil
+# import traceback
+# import gradio as gr
+# from pathlib import Path
+# from histopath.agent import A1
+# from dotenv import load_dotenv
+
+# # Load environment variables
+# load_dotenv()
+
+# # Get passcode from environment
+# PASSCODE = os.getenv("GRADIO_PASSWORD")
+
+# # Initialize agent (will be created after passcode validation)
+# agent = None
+
+
+# def check_for_output_files():
+#     """Check for all files in the output directory and return their paths."""
+#     output_dir = Path("./output")
+#     if not output_dir.exists():
+#         return [], []
+    
+#     image_extensions = {".png", ".jpg", ".jpeg", ".svg", ".tif", ".tiff"}
+#     data_extensions = {".csv", ".txt", ".json", ".npy"}
+    
+#     images = []
+#     data_files = []
+    
+#     for file in output_dir.iterdir():
+#         if file.is_file():
+#             if file.suffix.lower() in image_extensions:
+#                 images.append(str(file))
+#             elif file.suffix.lower() in data_extensions:
+#                 data_files.append(str(file))
+    
+#     return images, data_files
+
+
+# def preview_uploaded_file(uploaded_file):
+#     """Preview the uploaded file - show image or file info."""
+#     if uploaded_file is None:
+#         return None, None, "No file uploaded"
+    
+#     file_path = Path(uploaded_file.name)
+#     file_ext = file_path.suffix.lower()
+    
+#     image_extensions = {".png", ".jpg", ".jpeg", ".svg", ".tif", ".tiff", ".svs"}
+    
+#     if file_ext in image_extensions:
+#         # Show image preview
+#         return uploaded_file.name, None, f"📷 Previewing: {file_path.name}"
+#     else:
+#         # Show file info
+#         file_size = Path(uploaded_file.name).stat().st_size / 1024  # KB
+#         return None, uploaded_file.name, f"📄 File: {file_path.name} ({file_size:.1f} KB)"
+
+
+# def parse_agent_output(output):
+#     """Parse agent output to extract code blocks, observations, and regular text."""
+#     # Strip out the message divider bars
+#     output = re.sub(r'={30,}\s*(Human|Ai)\s+Message\s*={30,}', '', output)
+#     output = output.strip()
+    
+#     parsed = {
+#         "type": "text",
+#         "content": output,
+#         "code": None,
+#         "observation": None,
+#         "thinking": None
+#     }
+    
+#     # Check for code execution block
+#     execute_match = re.search(r'<execute>(.*?)</execute>', output, re.DOTALL)
+#     if execute_match:
+#         parsed["type"] = "code"
+#         parsed["code"] = execute_match.group(1).strip()
+#         # Extract text before the code block (thinking/explanation)
+#         text_before = output[:execute_match.start()].strip()
+#         # Remove any think tags but keep the content
+#         text_before = re.sub(r'<think>(.*?)</think>', r'\1', text_before, flags=re.DOTALL)
+#         text_before = re.sub(r'={30,}.*?={30,}', '', text_before).strip()
+#         parsed["thinking"] = text_before if text_before else None
+#         return parsed
+    
+#     # Check for observation block
+#     observation_match = re.search(r'<observation>(.*?)</observation>', output, re.DOTALL)
+#     if observation_match:
+#         parsed["type"] = "observation"
+#         parsed["observation"] = observation_match.group(1).strip()
+#         # Extract text before observation if any
+#         text_before = output[:observation_match.start()].strip()
+#         text_before = re.sub(r'<think>(.*?)</think>', r'\1', text_before, flags=re.DOTALL)
+#         text_before = re.sub(r'={30,}.*?={30,}', '', text_before).strip()
+#         parsed["thinking"] = text_before if text_before else None
+#         return parsed
+    
+#     # Check for solution block
+#     solution_match = re.search(r'<solution>(.*?)</solution>', output, re.DOTALL)
+#     if solution_match:
+#         parsed["type"] = "solution"
+#         parsed["content"] = solution_match.group(1).strip()
+#         # Get thinking before solution
+#         text_before = output[:solution_match.start()].strip()
+#         text_before = re.sub(r'<think>(.*?)</think>', r'\1', text_before, flags=re.DOTALL)
+#         text_before = re.sub(r'={30,}.*?={30,}', '', text_before).strip()
+#         parsed["thinking"] = text_before if text_before else None
+#         return parsed
+    
+#     # Clean up any remaining tags for display
+#     cleaned = re.sub(r'<think>(.*?)</think>', r'\1', output, flags=re.DOTALL)
+#     cleaned = re.sub(r'={30,}.*?={30,}', '', cleaned).strip()
+#     parsed["content"] = cleaned
+    
+#     return parsed
+
+
+# def format_message_for_display(parsed_output):
+#     """Format parsed output into a readable message for the chatbot."""
+#     msg_parts = []
+    
+#     # Add thinking/explanation text first if present
+#     if parsed_output.get("thinking"):
+#         msg_parts.append(parsed_output["thinking"])
+    
+#     if parsed_output["type"] == "code":
+#         # Add separator if there was thinking text
+#         if parsed_output.get("thinking"):
+#             msg_parts.append("\n---\n")
+        
+#         msg_parts.append("### 💻 Executing Code\n")
+#         msg_parts.append(f"```python\n{parsed_output['code']}\n```")
+        
+#     elif parsed_output["type"] == "observation":
+#         # Add separator if there was thinking text
+#         if parsed_output.get("thinking"):
+#             msg_parts.append("\n---\n")
+        
+#         msg_parts.append("### 📊 Observation\n")
+#         msg_parts.append(f"```\n{parsed_output['observation']}\n```")
+        
+#     elif parsed_output["type"] == "solution":
+#         # Add separator if there was thinking text
+#         if parsed_output.get("thinking"):
+#             msg_parts.append("\n---\n")
+        
+#         msg_parts.append("### ✅ Solution\n")
+#         msg_parts.append(parsed_output['content'])
+        
+#     else:
+#         # For regular text, just add the content if thinking wasn't already set
+#         if not parsed_output.get("thinking"):
+#             msg_parts.append(parsed_output["content"])
+    
+#     return "\n\n".join(msg_parts)
+
+
+# def process_agent_response(prompt, uploaded_file, chatbot_history):
+#     """Process the agent response and update chatbot."""
+#     global agent
+    
+#     if agent is None:
+#         chatbot_history.append({
+#             "role": "assistant",
+#             "content": "⚠️ Please enter the passcode first to initialize the agent."
+#         })
+#         yield chatbot_history, None, None, None, None, "⚠️ Agent not initialized"
+#         return
+    
+#     if not prompt.strip() and uploaded_file is None:
+#         chatbot_history.append({
+#             "role": "assistant",
+#             "content": "⚠️ Please provide a prompt or upload a file."
+#         })
+#         yield chatbot_history, None, None, None, None, "⚠️ No input provided"
+#         return
+    
+#     # Handle file upload
+#     file_path = None
+#     file_info = ""
+#     if uploaded_file is not None:
+#         try:
+#             # Create data directory if it doesn't exist
+#             data_dir = Path("./data")
+#             data_dir.mkdir(exist_ok=True)
+            
+#             # Copy uploaded file to data directory
+#             file_name = Path(uploaded_file.name).name
+#             file_path = data_dir / file_name
+#             shutil.copy(uploaded_file.name, file_path)
+            
+#             file_info = f"\n\n📎 **Uploaded file:** `{file_path}`\n"
+            
+#             # Augment prompt with file path
+#             if prompt.strip():
+#                 prompt = f"{prompt}\n\nUploaded file path: {file_path}"
+#             else:
+#                 prompt = f"I have uploaded a file at: {file_path}. Please analyze it."
+                
+#         except Exception as e:
+#             error_msg = f"❌ Error handling file upload: {str(e)}"
+#             chatbot_history.append({
+#                 "role": "assistant",
+#                 "content": error_msg
+#             })
+#             yield chatbot_history, None, None, None, None, error_msg
+#             return
+    
+#     # Add user message to chat
+#     user_message = prompt if not file_info else f"{prompt}{file_info}"
+#     chatbot_history.append({"role": "user", "content": user_message})
+#     yield chatbot_history, None, None, None, None, "🔄 Processing..."
+    
+#     try:
+#         # Stream agent responses
+#         step_count = 0
+#         for step in agent.go_stream(prompt):
+#             step_count += 1
+#             output = step.get("output", "")
+            
+#             if output:
+#                 # Parse the output
+#                 parsed = parse_agent_output(output)
+                
+#                 # Add thinking text as separate message if present
+#                 if parsed.get("thinking"):
+#                     chatbot_history.append({
+#                         "role": "assistant",
+#                         "content": parsed["thinking"]
+#                     })
+                
+#                 # Add the block (code/observation/solution) as separate message if present
+#                 if parsed["type"] == "code" and parsed["code"]:
+#                     chatbot_history.append({
+#                         "role": "assistant",
+#                         "content": f"### 💻 Executing Code\n\n```python\n{parsed['code']}\n```"
+#                     })
+#                 elif parsed["type"] == "observation" and parsed["observation"]:
+#                     chatbot_history.append({
+#                         "role": "assistant",
+#                         "content": f"### 📊 Observation\n\n```\n{parsed['observation']}\n```"
+#                     })
+#                 elif parsed["type"] == "solution":
+#                     chatbot_history.append({
+#                         "role": "assistant",
+#                         "content": f"### ✅ Solution\n\n{parsed['content']}"
+#                     })
+#                 elif parsed["type"] == "text" and parsed["content"]:
+#                     # Only add if we haven't already added it as thinking
+#                     if not parsed.get("thinking"):
+#                         chatbot_history.append({
+#                             "role": "assistant",
+#                             "content": parsed["content"]
+#                         })
+                
+#                 # Check for output files after each step
+#                 images, data_files = check_for_output_files()
+                
+#                 # Create status message
+#                 status = f"🔄 Step {step_count}"
+#                 if parsed["type"] == "code":
+#                     status += " - Executing code..."
+#                 elif parsed["type"] == "observation":
+#                     status += " - Processing results..."
+#                 elif parsed["type"] == "solution":
+#                     status += " - Finalizing solution..."
+                
+#                 yield (
+#                     chatbot_history,
+#                     images if images else None,
+#                     data_files if data_files else None,
+#                     None,
+#                     None,
+#                     status
+#                 )
+        
+#         # Final check for files
+#         final_images, final_data = check_for_output_files()
+        
+#         # Create download links message if files were generated
+#         if final_images or final_data:
+#             download_msg = "\n\n---\n\n### 📁 Generated Files Ready for Download\n\n"
+            
+#             if final_images:
+#                 download_msg += f"**🖼️ Images ({len(final_images)})** - Available in the **Images** tab →\n"
+#                 for img_path in final_images:
+#                     img_name = Path(img_path).name
+#                     download_msg += f"- `{img_name}`\n"
+#                 download_msg += "\n"
+            
+#             if final_data:
+#                 download_msg += f"**📄 Data Files ({len(final_data)})** - Available in the **Data** tab →\n"
+#                 for data_path in final_data:
+#                     data_name = Path(data_path).name
+#                     download_msg += f"- `{data_name}`\n"
+            
+#             download_msg += "\n*Click the download button on each file in the respective tabs above.*"
+            
+#             # Add download message as separate bubble
+#             chatbot_history.append({
+#                 "role": "assistant",
+#                 "content": download_msg
+#             })
+        
+#         status = "✅ Complete"
+#         if final_images:
+#             status += f" | {len(final_images)} image(s)"
+#         if final_data:
+#             status += f" | {len(final_data)} data file(s)"
+            
+#         yield chatbot_history, final_images if final_images else None, final_data if final_data else None, None, None, status
+        
+#     except Exception as e:
+#         error_msg = f"❌ Error: {str(e)}\n\n```\n{traceback.format_exc()}\n```"
+#         chatbot_history.append({
+#             "role": "assistant",
+#             "content": error_msg
+#         })
+#         yield chatbot_history, None, None, None, None, "❌ Error occurred"
+
+
+# def validate_passcode(passcode):
+#     """Validate the passcode and initialize the agent."""
+#     global agent
+    
+#     if passcode == PASSCODE:
+#         # Initialize agent
+#         try:
+#             agent = A1()
+#             return (
+#                 gr.update(visible=False),  # Hide passcode section
+#                 gr.update(visible=True),   # Show main interface
+#                 "✅ Access granted! Agent initialized and ready."
+#             )
+#         except Exception as e:
+#             error_trace = traceback.format_exc()
+#             return (
+#                 gr.update(visible=True),
+#                 gr.update(visible=False),
+#                 f"❌ Error initializing agent:\n{str(e)}\n\n{error_trace}"
+#             )
+#     else:
+#         return (
+#             gr.update(visible=True),
+#             gr.update(visible=False),
+#             "❌ Invalid passcode. Please try again."
+#         )
+
+
+# def clear_chat():
+#     """Clear the chat history and output files."""
+#     # Clean up output directory
+#     output_dir = Path("./output")
+#     if output_dir.exists():
+#         shutil.rmtree(output_dir)
+#     output_dir.mkdir(exist_ok=True)
+    
+#     # Clean up data directory
+#     data_dir = Path("./data")
+#     if data_dir.exists():
+#         for file in data_dir.iterdir():
+#             if file.is_file():
+#                 file.unlink()
+    
+#     return [], None, None, None, None, "🗑️ Chat cleared"
+
+
+# # Create Gradio interface with custom theme
+# custom_theme = gr.themes.Soft(
+#     primary_hue="blue",
+#     secondary_hue="slate",
+#     spacing_size="sm",
+#     radius_size="md",
+# ).set(
+#     button_primary_background_fill="*primary_500",
+#     button_primary_background_fill_hover="*primary_600",
+#     block_label_text_weight="600",
+#     block_title_text_weight="600",
+# )
+
+# with gr.Blocks(title="HistoPath Agent", theme=custom_theme, css="""
+#     .gradio-container {
+#         max-width: 100% !important;
+#     }
+#     .main-header {
+#         text-align: center;
+#         padding: 1.5rem 0;
+#         background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+#         color: white;
+#         border-radius: 8px;
+#         margin-bottom: 1.5rem;
+#     }
+#     .main-header h1 {
+#         margin: 0;
+#         font-size: 2.2rem;
+#         font-weight: 700;
+#     }
+#     .main-header p {
+#         margin: 0.5rem 0 0 0;
+#         opacity: 0.95;
+#         font-size: 1.1rem;
+#     }
+#     .file-upload-box .wrap {
+#         min-width: 0 !important;
+#     }
+#     .file-upload-box .file-name {
+#         word-break: break-word !important;
+#         white-space: normal !important;
+#         overflow-wrap: break-word !important;
+#     }
+#     .tab-nav {
+#         margin-bottom: 0.5rem;
+#     }
+#     /* Better styling for code and observation blocks */
+#     .message.bot pre {
+#         background-color: #f6f8fa !important;
+#         border: 1px solid #d0d7de !important;
+#         border-radius: 6px !important;
+#         padding: 12px !important;
+#         margin: 8px 0 !important;
+#     }
+#     .message.bot h3 {
+#         margin-top: 12px !important;
+#         margin-bottom: 8px !important;
+#         font-weight: 600 !important;
+#     }
+#     .message.bot hr {
+#         border: none !important;
+#         border-top: 2px solid #e1e4e8 !important;
+#         margin: 16px 0 !important;
+#     }
+# """) as demo:
+    
+#     # Header
+#     gr.HTML("""
+#         <div class="main-header">
+#             <h1>🔬 HistoPath Agent</h1>
+#             <p>AI-Powered Histopathology Analysis Assistant</p>
+#         </div>
+#     """)
+    
+#     # Passcode section
+#     with gr.Group(visible=True) as passcode_section:
+#         gr.Markdown("### 🔐 Authentication Required")
+        
+#         with gr.Row():
+#             passcode_input = gr.Textbox(
+#                 label="Passcode",
+#                 type="password",
+#                 placeholder="Enter your passcode...",
+#                 scale=3
+#             )
+#             passcode_btn = gr.Button("🔓 Unlock", variant="primary", scale=1, size="lg")
+        
+#         passcode_status = gr.Textbox(
+#             label="Status",
+#             interactive=False,
+#             lines=2
+#         )
+    
+#     # Main interface (hidden initially)
+#     with gr.Group(visible=False) as main_interface:
+#         with gr.Row(equal_height=True):
+#             # Left column - Chat interface
+#             with gr.Column(scale=3):
+#                 chatbot = gr.Chatbot(
+#                     label="💬 Conversation",
+#                     height=550,
+#                     # type="messages",
+#                     show_label=True,
+#                     render_markdown=True,
+#                 )
+                
+#                 # Input area
+#                 with gr.Row():
+#                     with gr.Column(scale=7):
+#                         prompt_input = gr.Textbox(
+#                             label="Your Query",
+#                             placeholder="E.g., 'Caption the uploaded whole slide image' or 'Segment cells using instanseg model'",
+#                             lines=2,
+#                             max_lines=5,
+#                             show_label=False,
+#                         )
+#                     with gr.Column(scale=3):
+#                         file_upload = gr.File(
+#                             label="📎 Upload File",
+#                             file_types=[".svs", ".png", ".jpg", ".jpeg", ".tif", ".tiff", ".csv", ".txt", ".json", ".npy"],
+#                             height=75,
+#                             elem_classes="file-upload-box",
+#                         )
+                
+#                 with gr.Row():
+#                     submit_btn = gr.Button("🚀 Submit", variant="primary", scale=3, size="lg")
+#                     clear_btn = gr.Button("🗑️ Clear", scale=1, size="lg", variant="secondary")
+                
+#                 status_text = gr.Textbox(
+#                     label="Status",
+#                     interactive=False,
+#                     value="Ready",
+#                     show_label=False,
+#                     container=False,
+#                 )
+            
+#             # Right column - Outputs
+#             with gr.Column(scale=2):
+#                 with gr.Tabs():
+#                     with gr.Tab("📥 Input"):
+#                         with gr.Column():
+#                             input_image_preview = gr.Image(
+#                                 label="Input Image",
+#                                 height=400,
+#                                 show_label=False,
+#                                 container=True,
+#                             )
+#                             input_file_preview = gr.File(
+#                                 label="Input File",
+#                                 interactive=False,
+#                                 height=100,
+#                                 show_label=False,
+#                                 container=True,
+#                             )
+#                             input_status = gr.Textbox(
+#                                 value="Upload a file to preview",
+#                                 show_label=False,
+#                                 interactive=False,
+#                                 container=False,
+#                             )
+                    
+#                     with gr.Tab("🖼️ Images"):
+#                         output_gallery = gr.Gallery(
+#                             label="Generated Visualizations",
+#                             columns=1,
+#                             height=600,
+#                             object_fit="contain",
+#                             show_label=False,
+#                             show_download_button=True,
+#                         )
+                    
+#                     with gr.Tab("📄 Data"):
+#                         data_files = gr.File(
+#                             label="Generated Data Files",
+#                             file_count="multiple",
+#                             interactive=False,
+#                             height=600,
+#                             show_label=False,
+#                         )
+    
+#     # Event handlers
+#     passcode_btn.click(
+#         fn=validate_passcode,
+#         inputs=[passcode_input],
+#         outputs=[passcode_section, main_interface, passcode_status]
+#     )
+    
+#     # File upload preview
+#     file_upload.change(
+#         fn=preview_uploaded_file,
+#         inputs=[file_upload],
+#         outputs=[input_image_preview, input_file_preview, input_status]
+#     )
+    
+#     submit_btn.click(
+#         fn=process_agent_response,
+#         inputs=[prompt_input, file_upload, chatbot],
+#         outputs=[chatbot, output_gallery, data_files, input_image_preview, input_file_preview, status_text]
+#     )
+    
+#     clear_btn.click(
+#         fn=clear_chat,
+#         outputs=[chatbot, output_gallery, data_files, input_image_preview, input_file_preview, status_text]
+#     )
+    
+#     # Allow enter key to submit
+#     prompt_input.submit(
+#         fn=process_agent_response,
+#         inputs=[prompt_input, file_upload, chatbot],
+#         outputs=[chatbot, output_gallery, data_files, input_image_preview, input_file_preview, status_text]
+#     )
+
+
+# if __name__ == "__main__":
+#     # Create necessary directories
+#     Path("./data").mkdir(exist_ok=True)
+#     Path("./output").mkdir(exist_ok=True)
+    
+#     print("=" * 60)
+#     print("🔬 HistoPath Agent - Gradio Interface")
+#     print("=" * 60)
+#     # print(f"Passcode: {PASSCODE}")
+#     print("Starting server...")
+#     print("=" * 60)
+    
+#     # Launch the app
+#     #demo.launch(
+#         # server_name="0.0.0.0",
+#         # server_port=7860,  # Let Gradio auto-pick an available port
+#         # share=True,
+#         # show_error=True,
+#     #)
+#     demo.launch()