Update app.py

app.py

@@ -33,7 +33,7 @@ from smolagents import (
     GoogleSearchTool,
     Tool,
 )
-from smolagents.agent_types import AgentText, AgentImage, AgentAudio
+from smolagents.agent_types import AgentText # Removed: AgentImage, AgentAudio
 from smolagents.gradio_ui import pull_messages_from_step, handle_agent_output_types
 
 
@@ -162,8 +162,8 @@ class ModelManager:
                 )
 
             raise ValueError(f"Invalid inference type: {chosen_inference}")
-        except Exception as e:
-            print(f"✗ Couldn't load model: {e}")
+        except (ValueError, RuntimeError) as e:  # More specific exceptions
+            print(f"Model loading failed: {e}")
             raise
 
 
@@ -191,7 +191,10 @@ class ToolRegistry:
             return Tool.from_space(
                 space_id="xkerser/FLUX.1-dev",
                 name="image_generator",
-                description="Generates high-quality AgentImage with text prompt (77 token limit).",
+                description=(
+                    "Generates high-quality AgentImage. "
+                    "with text prompt (77 token limit).",
+                ),
             )
         except Exception as e:
             print(f"✗ Couldn't initialize image generation tool: {e}")
@@ -251,43 +254,68 @@ def stream_to_gradio(
     reset_agent_memory: bool = False,
     additional_args: Optional[dict] = None,
 ):
-    """Runs an agent with the given task and streams messages as Gradio ChatMessages."""
-    for step_log in agent.run(
-        task, stream=True, reset=reset_agent_memory, additional_args=additional_args
-    ):
-        yield from pull_messages_from_step(step_log)
+    """Streams agent responses with improved status indicators."""
+    try:
+        # Initial processing indicator
+        yield gr.ChatMessage(role="assistant", content="⏳ Processing your request...")
+
+        # Track what we've yielded to replace the processing indicator
+        first_message_yielded = False
+
+        for step_log in agent.run(
+            task, stream=True, reset=reset_agent_memory, additional_args=additional_args
+        ):
+            messages = pull_messages_from_step(step_log)
+
+            for message in messages:
+                if not first_message_yielded:
+                    # Replace the initial "Processing" message
+                    first_message_yielded = True
+                    message.content = message.content.replace(
+                        "⏳ Processing your request...", ""
+                    )
 
-    # Process final answer : Use a more comprehensive media output
-    final_answer = step_log  # Last log is the run's final_answer
-    final_answer = handle_agent_output_types(final_answer)
+                # Add progressive status indicators for document processing steps
+                if "analyzing document" in step_log.get("status", ""):
+                    message.content = f"📄 **Document Analysis:** {message.content}"
+                elif "searching" in step_log.get("status", ""):
+                    message.content = f"🔍 **Search:** {message.content}"
 
-    if isinstance(final_answer, AgentText):
-        yield gr.ChatMessage(
-            role="assistant",
-            content=f"**Final answer:**\n{final_answer.to_string()}\n",
-        )
-    elif isinstance(final_answer, AgentImage):
-        yield gr.ChatMessage(
-            role="assistant",
-            content={"image": final_answer.to_string(), "type": "file"},
-        )  # Send as Gradio-compatible file object:
-    elif isinstance(final_answer, AgentAudio):
+                yield message
+
+        # Final answer with enhanced formatting
+        final_answer = handle_agent_output_types(step_log)
+
+        if isinstance(final_answer, AgentText):
+            yield gr.ChatMessage(
+                role="assistant",
+                content=f"✅ **Final Answer:**\n\n{final_answer.to_string()}",
+            )
+        else:
+            yield gr.ChatMessage(
+                role="assistant", content=f"✅ **Final Answer:** {str(final_answer)}"
+            )
+
+    except Exception as e:
         yield gr.ChatMessage(
             role="assistant",
-            content={"audio": final_answer.to_string(), "type": "file"},
-        )  # Send as Gradio-compatible file object
-    else:
-        yield gr.ChatMessage(
-            role="assistant", content=f"**Final answer:** {str(final_answer)}"
+            content=f"❌ **Error:** {str(e)}\n\nPlease try again with a different query.",
         )
 
 
 class GradioUI:
-    """A one-line interface to launch your agent in Gradio."""
-
-    def __init__(self, file_upload_folder: str | None = None):
-        """Initialize the Gradio UI with optional file upload functionality."""
+    def __init__(self, file_upload_folder=None, max_queue_size=50):
+        # Initialize all attributes here
         self.file_upload_folder = file_upload_folder
+        self.max_queue_size = max_queue_size
+        self.text_input = None
+        self.submit_btn = None
+        self.stop_btn = None
+        self.clear_btn = None
+        self.status = None
+        self.chatbot = None
+        self.session_state = None
+        self.job = None
 
         if self.file_upload_folder is not None:
             os.makedirs(file_upload_folder, exist_ok=True)
@@ -377,7 +405,7 @@ class GradioUI:
 
     def log_user_message(self, text_input, file_uploads_log):
         """Process user message and handle file references."""
-        # Clean the user input using the TextCleanerTool
+
         cleaned_message = clean(
             text_input,
             fix_unicode=True,
@@ -431,8 +459,7 @@ class GradioUI:
         if platform:
             if platform in ['"android"', '"ios"']:
                 return "Mobile"
-            elif platform in ['"windows"', '"macos"', '"linux"']:
-                return "Desktop"
+            return "Desktop"
 
         # Default case if no clear indicators
         return "Desktop"
@@ -579,56 +606,75 @@ class GradioUI:
 
             return simple_demo
 
-    def _connect_event_handlers(
-        self,
-        text_input,
-        launch_research_btn,
-        file_uploads_log,
-        stored_messages,
-        chatbot,
-        session_state,
-    ):
-        """Connect the event handlers for input elements."""
-        # Connect text input submit event
-        text_input.submit(
-            self.log_user_message,
-            [text_input, file_uploads_log],
-            [stored_messages, text_input, launch_research_btn],
-        ).then(
-            self.interact_with_agent,
-            [stored_messages, chatbot, session_state],
-            [chatbot],
-        ).then(
-            lambda: (
-                gr.Textbox(
-                    interactive=True,
-                    placeholder="Enter your prompt here and press the button",
-                ),
-                gr.Button(interactive=True),
-            ),
-            None,
-            [text_input, launch_research_btn],
+    def _create_common_ui_elements(self):
+        """Create common UI elements with control buttons."""
+        with gr.Group():
+            self.text_input = gr.Textbox(
+                lines=3,
+                label="Your request",
+                placeholder="Enter your question about the documents...",
+                elem_classes=["prompt-box"],
+            )
+
+            with gr.Row():
+                self.submit_btn = gr.Button("Run", variant="primary")
+                self.stop_btn = gr.Button("Stop Generation", variant="stop")
+                self.clear_btn = gr.Button("Clear Chat", variant="secondary")
+
+            # Status indicator for document processing
+            self.status = gr.Textbox(
+                "", label="Status", interactive=False, visible=True
+            )
+
+    def _connect_event_handlers(self):
+        """Connect event handlers with stop functionality."""
+        # Define the job handler for stopping generation
+        self.job = None
+
+        def start_processing(prompt, chat_history):
+            # Update UI for processing state
+            self.status.update(value="⏳ Processing request...", visible=True)
+            self.submit_btn.update(interactive=False)
+            self.stop_btn.update(visible=True)
+            return prompt, chat_history
+
+        def stop_generation():
+            if self.job:
+                self.job.cancel()
+            self.status.update(value="✋ Generation stopped by user", visible=True)
+            self.submit_btn.update(interactive=True)
+            self.stop_btn.update(visible=False)
+
+        def clear_chat():
+            return [], gr.Textbox(interactive=True), gr.Button(interactive=True), ""
+
+        # Connect event handlers
+        process_event = (
+            self.text_input.submit(
+                start_processing,
+                [self.text_input, self.chatbot],
+                [self.text_input, self.chatbot],
+            )
+            .then(
+                self._interact_with_agent,
+                [self.text_input, self.chatbot, self.session_state],
+                [self.chatbot],
+            )
+            .then(
+                lambda: (gr.Textbox(interactive=True), gr.Button(interactive=True), ""),
+                None,
+                [self.text_input, self.submit_btn, self.status],
+            )
         )
 
-        # Connect button click event
-        launch_research_btn.click(
-            self.log_user_message,
-            [text_input, file_uploads_log],
-            [stored_messages, text_input, launch_research_btn],
-        ).then(
-            self.interact_with_agent,
-            [stored_messages, chatbot, session_state],
-            [chatbot],
-        ).then(
-            lambda: (
-                gr.Textbox(
-                    interactive=True,
-                    placeholder="Enter your prompt here and press the button",
-                ),
-                gr.Button(interactive=True),
-            ),
-            None,
-            [text_input, launch_research_btn],
+        # Store the job for cancellation
+        self.job = process_event
+
+        # Connect stop and clear buttons
+        self.stop_btn.click(stop_generation)
+        self.clear_btn.click(
+            clear_chat,
+            outputs=[self.chatbot, self.text_input, self.submit_btn, self.status],
         )