Update app.py

app.py

@@ -33,7 +33,7 @@ from smolagents import (
     GoogleSearchTool,
     Tool,
 )
-from smolagents.agent_types import AgentText # Removed: AgentImage, AgentAudio
+from smolagents.agent_types import AgentText  # AgentImage, AgentAudio
 from smolagents.gradio_ui import pull_messages_from_step, handle_agent_output_types
 
 
@@ -193,7 +193,7 @@ class ToolRegistry:
                 name="image_generator",
                 description=(
                     "Generates high-quality AgentImage. "
-                    "with text prompt (77 token limit).",
+                    "with text prompt (77 token limit)."
                 ),
             )
         except Exception as e:
@@ -202,9 +202,9 @@ class ToolRegistry:
 
     @staticmethod
     def load_clean_text_tool():
-        """Initialize and return image generation tools."""
+        """Initialize and return text cleaning tool."""
         try:
-            return TextCleanerTool
+            return TextCleanerTool()
         except Exception as e:
             print(f"✗ Couldn't initialize clean text tool: {e}")
             raise
@@ -227,15 +227,14 @@ def create_agent():
     image_generator = ToolRegistry.load_image_generation_tools()
     clean_text = TextCleanerTool()  # Instantiate TextCleanerTool
 
-    # Combine all tools into a single list (not a tuple)
-    all_tools = [visualizer] + web_tools + [image_generator] + [clean_text]
+    # Combine all tools into a single list
+    all_tools = [visualizer] + web_tools + [image_generator, clean_text]
 
     # Validate tools before creating agent
     for tool in all_tools:
         if not isinstance(tool, Tool):
             raise ValueError(
-                "Invalid tool type: "
-                f"{type(tool)}. All tools must be instances of Tool class."
+                f"Invalid tool type: {type(tool)}. All tools must be instances of Tool class."
             )
 
     return CodeAgent(
@@ -626,57 +625,55 @@ class GradioUI:
                 "", label="Status", interactive=False, visible=True
             )
 
-    def _connect_event_handlers(self):
-        """Connect event handlers with stop functionality."""
+    def _connect_event_handlers(
+        self,
+        text_input,
+        launch_research_btn,
+        file_uploads_log,
+        stored_messages,
+        chatbot,
+        session_state,
+    ):
+        """Connect event handlers with appropriate parameters."""
         # 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)
+            # We'll use the passed components directly rather than self.status
             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 text input submission
+        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],
         )
 
-        # 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],
+        # Connect button click
+        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],
         )
 
+        # Store the job for cancellation if needed
+        self.job = None  # This would need to be assigned to an actual event
+
 
 def main():
     """Main entry point for the application."""