Thread safe browsing

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 import requests
 import pandas as pd
 import datasets
-from mini_agents import master_agent
+from mini_agents import MasterAgentWrapper
 from utils import get_full_file_path
 from smolagents.memory import ActionStep, PlanningStep, TaskStep, SystemPromptStep, FinalAnswerStep
 from typing import Optional
@@ -43,43 +43,52 @@ columns = [
 
 class BasicAgent:
     def __init__(self):
-        self.agent = master_agent
+        self.agent = MasterAgentWrapper()  # This is now the MasterAgentWrapper instance
         print("Master Agent initialized.")
-    def __call__(self, question: str, task_id: str, df_agent_steps: pd.DataFrame) -> str:
+        
+    def __call__(self, question: str, task_id: str, df_agent_steps: pd.DataFrame) -> tuple[str, pd.DataFrame]:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
-        fixed_answer = self.agent.run(question)
-        all_steps = self.agent.memory.get_full_steps()
-        for step in all_steps:
-            if isinstance(step, ActionStep):
-                step_class = "ActionStep"
-            elif isinstance(step, PlanningStep):
-                step_class = "PlanningStep"
-            elif isinstance(step, TaskStep):
-                step_class = "TaskStep"
-            elif isinstance(step, SystemPromptStep):
-                step_class = "SystemPromptStep"
-            elif isinstance(step, FinalAnswerStep):
-                step_class = "FinalAnswerStep"
-            else:
-                step_class = "UnknownStep"
+        try:
+            # Use the wrapper's run method which handles browser tools safely
+            fixed_answer = self.agent.run(question)
+            
+            # Log steps
+            all_steps = self.agent.master_agent.memory.get_full_steps()
+            for step in all_steps:
+                if isinstance(step, ActionStep):
+                    step_class = "ActionStep"
+                elif isinstance(step, PlanningStep):
+                    step_class = "PlanningStep"
+                elif isinstance(step, TaskStep):
+                    step_class = "TaskStep"
+                elif isinstance(step, SystemPromptStep):
+                    step_class = "SystemPromptStep"
+                elif isinstance(step, FinalAnswerStep):
+                    step_class = "FinalAnswerStep"
+                else:
+                    step_class = "UnknownStep"
 
-            step_dict = step.dict()
-            df_agent_steps.loc[len(df_agent_steps)] = None
-            df_agent_steps.at[len(df_agent_steps), 'task_id'] = task_id
-            df_agent_steps.at[len(df_agent_steps), 'step_class'] = step_class
-            for key, value in step_dict.items():
-                df_agent_steps.at[len(df_agent_steps), key] = value
+                step_dict = step.dict()
+                df_agent_steps.loc[len(df_agent_steps)] = None
+                df_agent_steps.at[len(df_agent_steps), 'task_id'] = task_id
+                df_agent_steps.at[len(df_agent_steps), 'step_class'] = step_class
+                for key, value in step_dict.items():
+                    df_agent_steps.at[len(df_agent_steps), key] = value
             
-        print(f"Agent returning fixed answer: {fixed_answer}")
-        return fixed_answer, df_agent_steps
+            print(f"Agent returning fixed answer: {fixed_answer}")
+            return fixed_answer, df_agent_steps
+            
+        except Exception as e:
+            print(f"Error in agent execution: {e}")
+            raise
 
 def check_required_env_vars() -> tuple[bool, Optional[str]]:
     """Check if required environment variables are set"""
     missing_vars = []
     
     # Check HF_TOKEN
-    if not os.getenv("HF_TOKEN"):
-        missing_vars.append("HF_TOKEN")
+    if not os.getenv("HUGGINGFACE_API_KEY"):
+        missing_vars.append("HUGGINGFACE_API_KEY")
     
     # Check SPACE_ID (only warn, not required)
     if not os.getenv("SPACE_ID"):
@@ -117,7 +126,7 @@ def save_dataset_to_hub(df: pd.DataFrame, dataset_name: str) -> tuple[bool, str]
         dataset.push_to_hub(
             dataset_name,
             private=True,
-            token=os.getenv("HF_TOKEN")
+            token=os.getenv("HUGGINGFACE_API_KEY")
         )
         
         return True, f"Successfully saved {len(df)} steps to {dataset_name}"

browser.py

@@ -1,16 +1,75 @@
 import subprocess
-
-subprocess.run(["bash", "scripts.sh"])
-
+from contextlib import contextmanager
+from typing import List, Optional
+from threading import Lock
 from smolagents.tools import Tool
 from langchain_community.tools.playwright.utils import (
     create_async_playwright_browser,
     create_sync_playwright_browser
-
 )
 from langchain_community.agent_toolkits import PlayWrightBrowserToolkit
 
-async_browser = create_async_playwright_browser()
-sync_browser = create_sync_playwright_browser()
-browser_toolkit = PlayWrightBrowserToolkit.from_browser(async_browser=async_browser, sync_browser=sync_browser)
-browser_tools = [Tool.from_langchain(tool) for tool in browser_toolkit.get_tools()]
+class BrowserManager:
+    _instance = None
+    _lock = Lock()
+    _browser_tools: Optional[List[Tool]] = None
+    
+    def __new__(cls):
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super(BrowserManager, cls).__new__(cls)
+        return cls._instance
+    
+    def __init__(self):
+        if not hasattr(self, 'initialized'):
+            # Run setup script
+            subprocess.run(["bash", "scripts.sh"])
+            self.initialized = True
+    
+    @contextmanager
+    def get_browser_tools(self):
+        """Get browser tools in a context that ensures proper cleanup"""
+        try:
+            if self._browser_tools is None:
+                with self._lock:
+                    if self._browser_tools is None:
+                        # Create browsers in the current context
+                        async_browser = create_async_playwright_browser()
+                        sync_browser = create_sync_playwright_browser()
+                        
+                        # Create toolkit and tools
+                        browser_toolkit = PlayWrightBrowserToolkit.from_browser(
+                            async_browser=async_browser,
+                            sync_browser=sync_browser
+                        )
+                        self._browser_tools = [
+                            Tool.from_langchain(tool) 
+                            for tool in browser_toolkit.get_tools()
+                        ]
+            
+            yield self._browser_tools
+            
+        except Exception as e:
+            print(f"Error in browser context: {e}")
+            # Reset tools on error
+            self._browser_tools = None
+            raise
+        finally:
+            # Cleanup if needed
+            if self._browser_tools:
+                for tool in self._browser_tools:
+                    if hasattr(tool, 'browser'):
+                        try:
+                            tool.browser.close()
+                        except:
+                            pass
+                self._browser_tools = None
+
+# Create a singleton instance
+browser_manager = BrowserManager()
+
+# For backward compatibility, but prefer using browser_manager.get_browser_tools()
+def get_browser_tools():
+    """Get browser tools (use with context manager)"""
+    return browser_manager.get_browser_tools()

mini_agents.py

@@ -5,10 +5,12 @@ from tools import to_dataframe, to_json, get_dataframe_data, get_dataframe_colum
 from vlm_tools import image_processing, object_detection_tool, ocr_scan_tool, extract_images_from_video, get_image_from_file_path, get_video_from_file_path
 from audio_tools import transcribe_audio_tool, get_audio_from_file_path, noise_reduction, audio_segmentation, speaker_diarization
 from community_tools import community_tools, get_youtube_transcript_from_url
-from browser import browser_tools
+from browser import browser_manager
 import os
 import logging
 import yaml
+from typing import List
+from smolagents.tools import Tool
 
 logging.basicConfig(level=logging.DEBUG)
 
@@ -141,19 +143,58 @@ master_model = InferenceClientModel(
     token=os.getenv("HUGGINGFACE_API_KEY")
 )
 
-master_agent = CodeAgent(
-    model=master_model,
-    managed_agents=[audio_agent, vlm_agent, arithmetic_agent, pandas_agent],
-    tools=[sort_list, get_youtube_transcript_from_url, read_python_file_from_path, *community_tools, *browser_tools, tavily_search_tool],
-    add_base_tools=True,
-    max_steps=20,
-    additional_authorized_imports=AUTHORIZED_IMPORTS,
-    verbosity_level=logging.INFO,
-    planning_interval=4,
-    prompt_templates=PROMPT_TEMPLATE["master_agent"],
-    name="master_agent",
-    description="This agent is responsible for managing audio, vlm, arithmetic and pandas agents."
-)
+class MasterAgentWrapper:
+    """Wrapper class to manage master agent with thread-safe browser tools"""
+    def __init__(self):
+        self.base_tools = [
+            sort_list, 
+            get_youtube_transcript_from_url, 
+            read_python_file_from_path, 
+            *community_tools, 
+            tavily_search_tool
+        ]
+        
+        self.master_agent = CodeAgent(
+            model=master_model,
+            managed_agents=[audio_agent, vlm_agent, arithmetic_agent, pandas_agent],
+            tools=self.base_tools,  # Initialize without browser tools
+            add_base_tools=True,
+            max_steps=20,
+            additional_authorized_imports=AUTHORIZED_IMPORTS,
+            verbosity_level=logging.INFO,
+            planning_interval=4,
+            prompt_templates=PROMPT_TEMPLATE["master_agent"],
+            name="master_agent",
+            description="This agent is responsible for managing audio, vlm, arithmetic and pandas agents."
+        )
+    
+    def run(self, question: str) -> str:
+        """Run the agent with thread-safe browser tools"""
+        try:
+            # Get browser tools in the correct context
+            with browser_manager.get_browser_tools() as browser_tools:
+                # Temporarily add browser tools for this run
+                original_tools = self.master_agent.tools
+                self.master_agent.tools = original_tools + browser_tools
+                
+                try:
+                    # Run the agent
+                    result = self.master_agent.run(question)
+                    return result
+                finally:
+                    # Restore original tools
+                    self.master_agent.tools = original_tools
+                    
+        except Exception as e:
+            logging.error(f"Error in master agent run: {e}")
+            raise
+
+# Create the wrapped master agent
+master_agent = MasterAgentWrapper()
+
+# For backward compatibility
+def run_master_agent(question: str) -> str:
+    return master_agent.run(question)
 
 #TESTING 5