Update app.py

app.py

@@ -74,11 +74,11 @@ AUTHORIZED_IMPORTS = [
     "time",  # Measuring time
     "tempfile",  # Creating temporary files and directories
     # Data Visualization (if needed) - Consider security implications carefully
-    "matplotlib",  # Plotting library (basic charts) - Consider security implications
-    "seaborn",  # Statistical data visualization (more advanced) - Consider security implications
-    # Web Scraping (more specific/controlled) - Consider ethical implications carefully
+    "matplotlib",  # Plotting library (basic charts)
+    "seaborn",  # Statistical data visualization (more advanced)
+    # Web Scraping (more specific/controlled) - Consider ethical implications
     "lxml",  # Faster XML/HTML processing (alternative to bs4)
-    "selenium",  # Automated browser control (for dynamic websites) - HIGH SECURITY RISK
+    "selenium",  # Automated browser control (for dynamic websites)
     # Database interaction (if needed) - Handle credentials securely!
     "sqlite3",  # SQLite database access
     # "psycopg2",         # PostgreSQL adapter if needed
@@ -89,7 +89,10 @@ AUTHORIZED_IMPORTS = [
 ]
 
 
-USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0"
+USER_AGENT = (
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
+    "(KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0"
+)
 BROWSER_CONFIG = {
     "viewport_size": 1024 * 5,
     "downloads_folder": "downloads_folder",
@@ -215,7 +218,8 @@ def create_agent():
     # Initialize model
     model = LiteLLMModel(
         custom_role_conversions=CUSTOM_ROLE_CONVERSIONS,
-        model_id="openrouter/deepseek/deepseek-chat-v3-0324:free",  # currently serving:
+        # Currently serving:
+        model_id="openrouter/deepseek/deepseek-chat-v3-0324:free",
     )  # DEEPSEEK = openrouter/perplexity/r1-1776 <--- boss model
 
     # Initialize tools
@@ -234,7 +238,8 @@ def create_agent():
     for tool in all_tools:
         if not isinstance(tool, Tool):
             raise ValueError(
-                f"Invalid tool type: {type(tool)}. All tools must be instances of Tool class."
+                f"Invalid tool type: {type(tool)}. "
+                f"All tools must be instances of Tool class."
             )
 
     return CodeAgent(
@@ -264,9 +269,9 @@ def stream_to_gradio(
         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:
+            # pull_messages_from_step is a generator function that yields messages
+            # We need to iterate through each yielded message
+            for message in pull_messages_from_step(step_log):
                 if not first_message_yielded:
                     # Replace the initial "Processing" message
                     first_message_yielded = True
@@ -274,10 +279,14 @@ def stream_to_gradio(
                         "⏳ Processing your request...", ""
                     )
 
-                # Add progressive status indicators for document processing steps
-                if "analyzing document" in step_log.get("status", ""):
+                # Check message content for document analysis or search references
+                content_lower = (
+                    message.content.lower() if hasattr(message, "content") else ""
+                )
+
+                if "document analysis" in content_lower:
                     message.content = f"📄 **Document Analysis:** {message.content}"
-                elif "searching" in step_log.get("status", ""):
+                elif "search" in content_lower:
                     message.content = f"🔍 **Search:** {message.content}"
 
                 yield message
@@ -298,7 +307,10 @@ def stream_to_gradio(
     except Exception as e:
         yield gr.ChatMessage(
             role="assistant",
-            content=f"❌ **Error:** {str(e)}\n\nPlease try again with a different query.",
+            content=(
+                f"❌ **Error:** {str(e)}\n\n"
+                f"Please try again with a different query."
+            ),
         )
 
 
@@ -367,9 +379,8 @@ class GradioUI:
 
         # Sanitize file name
         original_name = os.path.basename(file.name)
-        sanitized_name = re.sub(
-            r"[^\w\-.]", "_", original_name
-        )  # Replace invalid chars with underscores
+        # Replace invalid chars with underscores
+        sanitized_name = re.sub(r"[^\w\-.]", "_", original_name)
 
         # Ensure the extension correlates to the mime type
         type_to_ext = {}
@@ -419,12 +430,16 @@ class GradioUI:
             no_currency_symbols=False,
             no_punct=False,
             lang="en",
-        )  # Can change default behaviour by instantiating an object of TextCleanerTool class.
+        )  # Can change default behaviour with TextCleanerTool
 
         message = cleaned_message  # Use the cleaned message
 
         if file_uploads_log:
-            message += f"\nYou have been provided with these files, which might be helpful or not: {file_uploads_log}"  # Added file list
+            # Added file list to message
+            message += (
+                f"\nYou have been provided with these files, which might be "
+                f"helpful or not: {file_uploads_log}"
+            )
 
         return (
             message,
@@ -495,7 +510,10 @@ class GradioUI:
                         lines=3,
                         label="Your request",
                         container=False,
-                        placeholder="Enter your prompt here and press Shift+Enter or press the button",
+                        placeholder=(
+                            "Enter your prompt here and press Shift+Enter or "
+                            "press the button"
+                        ),
                     )
                     launch_research_btn = gr.Button("Run", variant="primary")
 
@@ -516,9 +534,11 @@ class GradioUI:
                 with gr.Row():
                     gr.HTML(
                         """
-                    <div style="display: flex; align-items: center; gap: 8px; font-family: system-ui, -apple-system, sans-serif;">
+                    <div style="display: flex; align-items: center; gap: 8px; 
+                    font-family: system-ui, -apple-system, sans-serif;">
                     <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png"
-                         style="width: 32px; height: 32px; object-fit: contain;" alt="logo">
+                         style="width: 32px; height: 32px; object-fit: contain;" 
+                         alt="logo">
                     <a target="_blank" href="https://github.com/huggingface/smolagents">
                         <b>huggingface/smolagents</b>
                     </a>
@@ -527,7 +547,8 @@ class GradioUI:
                     )
 
             # Add session state to store session-specific data
-            session_state = gr.State({})  # Initialize empty state for each session
+            # Initialize empty state for each session
+            session_state = gr.State({})
             stored_messages = gr.State([])
             if "file_uploads_log" not in locals():
                 file_uploads_log = gr.State([])