..

app.py

@@ -1,3 +1,639 @@
+# 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 - AGGRESSIVE FIX: Minimal yields."""
+#     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})
+    
+#     # CRITICAL FIX: Only yield once at the start to show user message
+#     yield chatbot_history, None, None, None, None, "🔄 Processing..."
+    
+#     try:
+#         # 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:
+#                 collected_outputs.append(output)
+        
+#         # 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
+#                 })
+        
+#         # CRITICAL FIX: Check files only ONCE after all processing
+#         images, data = check_for_output_files()
+        
+#         status_msg = f"✅ Complete ({step_count} steps)"
+#         if images:
+#             status_msg += f" | {len(images)} image(s)"
+#         if data:
+#             status_msg += f" | {len(data)} data file(s)"
+        
+#         # 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()
+#         error_msg = f"❌ **Error:** {str(e)}\n\n<details>\n<summary>Stack Trace</summary>\n\n```\n{error_trace}\n```\n</details>"
+        
+#         chatbot_history.append({
+#             "role": "assistant",
+#             "content": error_msg
+#         })
+        
+#         yield chatbot_history, None, None, None, None, f"❌ Error: {str(e)}"
+
+
+# def clear_chat():
+#     """Clear the chat history and outputs."""
+#     return [], None, None, None, None, "Ready"
+
+
+# def validate_passcode(input_passcode):
+#     """Validate the passcode and initialize the agent."""
+#     global agent
+    
+#     if input_passcode == PASSCODE:
+#         try:
+#             # Initialize the agent
+#             agent = A1(
+#                 path="./data",
+#                 llm="claude-sonnet-4-20250514",
+#                 use_tool_retriever=True,
+#                 timeout_seconds=600
+#             )
+#             return (
+#                 gr.update(visible=False),  # Hide passcode section
+#                 gr.update(visible=True),   # Show main interface
+#                 "✅ Access granted! Agent initialized successfully."
+#             )
+#         except Exception as e:
+#             error_msg = f"❌ Failed to initialize agent: {str(e)}"
+#             return (
+#                 gr.update(visible=True),
+#                 gr.update(visible=False),
+#                 error_msg
+#             )
+#     else:
+#         return (
+#             gr.update(visible=True),
+#             gr.update(visible=False),
+#             "❌ Invalid passcode. Please try again."
+#         )
+# # batched streaming instead
+# def process_agent_response_batched(prompt, uploaded_file, chatbot_history, batch_size=5):
+#     """Process agent response with BATCHED updates (every N steps)."""
+#     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:
+#             data_dir = Path("./data")
+#             data_dir.mkdir(exist_ok=True)
+            
+#             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"
+            
+#             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 with batching
+#         step_count = 0
+#         batch_count = 0
+        
+#         for step in agent.go_stream(prompt):
+#             step_count += 1
+#             output = step.get("output", "")
+            
+#             if output:
+#                 parsed = parse_agent_output(output)
+#                 formatted_message = format_message_for_display(parsed)
+                
+#                 # Update chatbot history
+#                 if chatbot_history and chatbot_history[-1]["role"] == "assistant":
+#                     chatbot_history[-1]["content"] = formatted_message
+#                 else:
+#                     chatbot_history.append({
+#                         "role": "assistant",
+#                         "content": formatted_message
+#                     })
+                
+#                 # Only yield every batch_size steps
+#                 if step_count % batch_size == 0:
+#                     batch_count += 1
+#                     yield chatbot_history, None, None, None, None, f"🔄 Step {step_count}..."
+        
+#         # Final yield with files
+#         images, data = check_for_output_files()
+        
+#         status_msg = f"✅ Complete ({step_count} steps)"
+#         if images:
+#             status_msg += f" | {len(images)} image(s)"
+#         if data:
+#             status_msg += f" | {len(data)} data file(s)"
+        
+#         yield chatbot_history, images, data, None, None, status_msg
+        
+#     except Exception as e:
+#         error_trace = traceback.format_exc()
+#         error_msg = f"❌ **Error:** {str(e)}\n\n<details>\n<summary>Stack Trace</summary>\n\n```\n{error_trace}\n```\n</details>"
+        
+#         chatbot_history.append({
+#             "role": "assistant",
+#             "content": error_msg
+#         })
+        
+#         yield chatbot_history, None, None, None, None, f"❌ Error: {str(e)}"
+
+
+# # Custom theme
+# custom_theme = gr.themes.Soft(
+#     primary_hue="indigo",
+#     secondary_hue="purple",
+#     neutral_hue="slate",
+#     font=["Inter", "system-ui", "sans-serif"],
+#     text_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",
+#                     type="messages",
+#                     height=550,
+#                     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("Starting server...")
+#     print("=" * 60)
+    
+#     # Launch the app
+#     demo.launch(show_api=False)
 import os
 import re
 import shutil
@@ -6,6 +642,7 @@ import gradio as gr
 from pathlib import Path
 from histopath.agent import A1
 from dotenv import load_dotenv
+from typing import List, Dict, Any, Optional, Tuple
 
 # Load environment variables
 load_dotenv()
@@ -157,8 +794,19 @@ def format_message_for_display(parsed_output):
     return "\n\n".join(msg_parts)
 
 
-def process_agent_response(prompt, uploaded_file, chatbot_history):
-    """Process the agent response and update chatbot - AGGRESSIVE FIX: Minimal yields."""
+# CRITICAL FIX: Wrap the main processing function to simplify type signatures
+def process_agent_response_wrapper(prompt: str, uploaded_file, chatbot_history: List[Dict[str, str]]):
+    """
+    Wrapper function with simplified type hints to avoid Gradio API introspection issues.
+    
+    Args:
+        prompt: User query string
+        uploaded_file: Uploaded file object (or None)
+        chatbot_history: List of message dictionaries
+        
+    Returns:
+        Tuple of (chatbot_history, images, data_files, preview_image, preview_file, status)
+    """
     global agent
     
     if agent is None:
@@ -223,56 +871,39 @@ def process_agent_response(prompt, uploaded_file, chatbot_history):
         for step in agent.go_stream(prompt):
             step_count += 1
             output = step.get("output", "")
-            
-            if output:
-                collected_outputs.append(output)
+            collected_outputs.append(output)
         
-        # CRITICAL FIX: Process ALL collected outputs at once
+        # Now process all collected outputs
         for output in collected_outputs:
+            if not output or output.strip() == "":
+                continue
+            
             parsed = parse_agent_output(output)
-            formatted_message = format_message_for_display(parsed)
+            formatted_msg = 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
+            if formatted_msg and formatted_msg.strip():
                 chatbot_history.append({
                     "role": "assistant",
-                    "content": formatted_message
+                    "content": formatted_msg
                 })
         
-        # CRITICAL FIX: Check files only ONCE after all processing
+        # CRITICAL: Single final yield with all results
         images, data = check_for_output_files()
+        final_status = f"✅ Complete! Processed {step_count} steps"
         
-        status_msg = f"✅ Complete ({step_count} steps)"
-        if images:
-            status_msg += f" | {len(images)} image(s)"
-        if data:
-            status_msg += f" | {len(data)} data file(s)"
-        
-        # CRITICAL FIX: Final single yield with all results
-        yield chatbot_history, images, data, None, None, status_msg
+        yield chatbot_history, images, data, None, None, final_status
         
     except Exception as e:
         error_trace = traceback.format_exc()
-        error_msg = f"❌ **Error:** {str(e)}\n\n<details>\n<summary>Stack Trace</summary>\n\n```\n{error_trace}\n```\n</details>"
-        
+        error_msg = f"❌ **Error occurred:**\n\n```\n{str(e)}\n\n{error_trace}\n```"
         chatbot_history.append({
             "role": "assistant",
             "content": error_msg
         })
-        
         yield chatbot_history, None, None, None, None, f"❌ Error: {str(e)}"
 
 
-def clear_chat():
-    """Clear the chat history and outputs."""
-    return [], None, None, None, None, "Ready"
-
-
-def validate_passcode(input_passcode):
+def validate_passcode(input_passcode: str) -> Tuple:
     """Validate the passcode and initialize the agent."""
     global agent
     
@@ -282,129 +913,40 @@ def validate_passcode(input_passcode):
             agent = A1(
                 path="./data",
                 llm="claude-sonnet-4-20250514",
+                source="Anthropic",
                 use_tool_retriever=True,
                 timeout_seconds=600
             )
+            
             return (
                 gr.update(visible=False),  # Hide passcode section
                 gr.update(visible=True),   # Show main interface
-                "✅ Access granted! Agent initialized successfully."
+                "✅ Authentication successful! Agent initialized."
             )
         except Exception as e:
-            error_msg = f"❌ Failed to initialize agent: {str(e)}"
             return (
-                gr.update(visible=True),
-                gr.update(visible=False),
-                error_msg
+                gr.update(visible=True),   # Keep passcode section visible
+                gr.update(visible=False),  # Keep main interface hidden
+                f"❌ Error initializing agent: {str(e)}"
             )
     else:
         return (
-            gr.update(visible=True),
-            gr.update(visible=False),
+            gr.update(visible=True),   # Keep passcode section visible
+            gr.update(visible=False),  # Keep main interface hidden
             "❌ Invalid passcode. Please try again."
         )
-# batched streaming instead
-def process_agent_response_batched(prompt, uploaded_file, chatbot_history, batch_size=5):
-    """Process agent response with BATCHED updates (every N steps)."""
-    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:
-            data_dir = Path("./data")
-            data_dir.mkdir(exist_ok=True)
-            
-            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"
-            
-            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..."
+
+
+def clear_chat():
+    """Clear the chat and reset outputs."""
+    # Clean output directory
+    output_dir = Path("./output")
+    if output_dir.exists():
+        for file in output_dir.iterdir():
+            if file.is_file():
+                file.unlink()
     
-    try:
-        # Stream with batching
-        step_count = 0
-        batch_count = 0
-        
-        for step in agent.go_stream(prompt):
-            step_count += 1
-            output = step.get("output", "")
-            
-            if output:
-                parsed = parse_agent_output(output)
-                formatted_message = format_message_for_display(parsed)
-                
-                # Update chatbot history
-                if chatbot_history and chatbot_history[-1]["role"] == "assistant":
-                    chatbot_history[-1]["content"] = formatted_message
-                else:
-                    chatbot_history.append({
-                        "role": "assistant",
-                        "content": formatted_message
-                    })
-                
-                # Only yield every batch_size steps
-                if step_count % batch_size == 0:
-                    batch_count += 1
-                    yield chatbot_history, None, None, None, None, f"🔄 Step {step_count}..."
-        
-        # Final yield with files
-        images, data = check_for_output_files()
-        
-        status_msg = f"✅ Complete ({step_count} steps)"
-        if images:
-            status_msg += f" | {len(images)} image(s)"
-        if data:
-            status_msg += f" | {len(data)} data file(s)"
-        
-        yield chatbot_history, images, data, None, None, status_msg
-        
-    except Exception as e:
-        error_trace = traceback.format_exc()
-        error_msg = f"❌ **Error:** {str(e)}\n\n<details>\n<summary>Stack Trace</summary>\n\n```\n{error_trace}\n```\n</details>"
-        
-        chatbot_history.append({
-            "role": "assistant",
-            "content": error_msg
-        })
-        
-        yield chatbot_history, None, None, None, None, f"❌ Error: {str(e)}"
+    return [], None, None, None, None, "Chat cleared"
 
 
 # Custom theme
@@ -603,7 +1145,7 @@ with gr.Blocks(title="HistoPath Agent", theme=custom_theme, css="""
     )
     
     submit_btn.click(
-        fn=process_agent_response,
+        fn=process_agent_response_wrapper,
         inputs=[prompt_input, file_upload, chatbot],
         outputs=[chatbot, output_gallery, data_files, input_image_preview, input_file_preview, status_text]
     )
@@ -615,7 +1157,7 @@ with gr.Blocks(title="HistoPath Agent", theme=custom_theme, css="""
     
     # Allow enter key to submit
     prompt_input.submit(
-        fn=process_agent_response,
+        fn=process_agent_response_wrapper,
         inputs=[prompt_input, file_upload, chatbot],
         outputs=[chatbot, output_gallery, data_files, input_image_preview, input_file_preview, status_text]
     )
@@ -632,5 +1174,7 @@ if __name__ == "__main__":
     print("Starting server...")
     print("=" * 60)
     
-    # Launch the app
-    demo.launch(show_api=False)
+    # Launch the app with API documentation completely disabled
+    demo.launch(
+        show_api=False,
+    )