import os import sys import importlib import inspect import traceback from typing import Dict, Any, List, Tuple import time from concurrent.futures import ThreadPoolExecutor, as_completed import multiprocessing def _get_optimal_workers(max_workers: int = None) -> tuple[int, str]: """ Intelligently determine optimal number of worker threads based on: 1. External parallelism (GNU parallel, SLURM, etc.) 2. Available CPU cores 3. User-specified max_workers Returns: (optimal_workers, reason): Number of workers and explanation """ cpu_count = multiprocessing.cpu_count() # Check for GNU parallel environment parallel_seq = os.environ.get('PARALLEL_SEQ') parallel_jobslot = os.environ.get('PARALLEL_JOBSLOT') # Check for SLURM scheduler slurm_ntasks = os.environ.get('SLURM_NTASKS') slurm_cpus_per_task = os.environ.get('SLURM_CPUS_PER_TASK') external_parallelism = None detection_method = None # Detect external parallelism if parallel_seq or parallel_jobslot: # GNU parallel is running # Estimate: typically parallel -j N means N parallel jobs # Conservative estimate: assume N = cpu_count / 2 external_parallelism = max(cpu_count // 2, 1) detection_method = "GNU parallel" elif slurm_ntasks: try: external_parallelism = int(slurm_ntasks) detection_method = "SLURM" except: pass # Calculate optimal workers if max_workers is not None: # User explicitly specified max_workers optimal = max_workers reason = f"user-specified: {max_workers} workers" elif external_parallelism: # External parallelism detected # Formula: optimal = max(1, cpu_count / external_parallelism) # But cap at 4 to avoid too few workers optimal = max(1, min(4, cpu_count // external_parallelism)) reason = f"{detection_method} detected, auto-adjusted to {optimal} workers (CPU={cpu_count}, external_jobs~{external_parallelism})" else: # No external parallelism, use moderate default optimal = min(4, max(2, cpu_count // 4)) reason = f"standalone mode: {optimal} workers (CPU={cpu_count})" return optimal, reason class Initializer: def __init__(self, enabled_tools: List[str] = [], tool_engine: List[str] = [], model_string: str = None, verbose: bool = False, vllm_config_path: str = None, base_url: str = None, check_model: bool = True, parallel_loading: bool = True, max_workers: int = None): """ Initialize the tool initializer with intelligent parallel loading. Args: enabled_tools: List of tool names to enable tool_engine: List of engine names corresponding to each tool model_string: Default model string verbose: Whether to print verbose output vllm_config_path: Path to vllm config base_url: Base URL for API check_model: Whether to check model availability parallel_loading: Whether to load tools in parallel (default: True) max_workers: Maximum number of parallel workers (default: None for auto-detect) If None, will intelligently detect based on: - External parallelism (GNU parallel, SLURM, etc.) - Available CPU cores """ self.toolbox_metadata = {} self.available_tools = [] self.enabled_tools = enabled_tools self.tool_engine = tool_engine self.load_all = self.enabled_tools == ["all"] self.model_string = model_string self.verbose = verbose self.vllm_server_process = None self.vllm_config_path = vllm_config_path self.base_url = base_url self.check_model = check_model self.parallel_loading = parallel_loading # Intelligently determine optimal workers optimal_workers, worker_reason = _get_optimal_workers(max_workers) self.max_workers = optimal_workers # Add tool instance cache - stores instantiated tools with their engines self.tool_instances_cache = {} print("\n==> Initializing agentflow...") print(f"Enabled tools: {self.enabled_tools} with {self.tool_engine}") print(f"LLM engine name: {self.model_string}") print(f"Parallel loading: {self.parallel_loading} ({worker_reason})") self._set_up_tools() # if vllm, set up the vllm server # if model_string.startswith("vllm-"): # self.setup_vllm_server() def get_project_root(self): current_dir = os.path.dirname(os.path.abspath(__file__)) while current_dir != '/': if os.path.exists(os.path.join(current_dir, 'agentflow')): return os.path.join(current_dir, 'agentflow') current_dir = os.path.dirname(current_dir) raise Exception("Could not find project root") def build_tool_name_mapping(self, tools_dir: str) -> Dict[str, Dict[str, str]]: """ Build a mapping dictionary by extracting TOOL_NAME from each tool file. Returns: Dict with two keys: - 'short_to_long': Maps short names (class names) to long names (external TOOL_NAME) - 'long_to_internal': Maps long names to internal class names and directory names """ short_to_long = {} # e.g., Base_Generator_Tool -> Generalist_Solution_Generator_Tool long_to_internal = {} # e.g., Generalist_Solution_Generator_Tool -> {class_name, dir_name} for root, dirs, files in os.walk(tools_dir): if 'tool.py' in files: dir_name = os.path.basename(root) tool_file_path = os.path.join(root, 'tool.py') try: # Read the tool.py file and extract TOOL_NAME with open(tool_file_path, 'r') as f: content = f.read() # Extract TOOL_NAME using simple string parsing external_tool_name = None for line in content.split('\n'): if line.strip().startswith('TOOL_NAME ='): # Extract the value between quotes external_tool_name = line.split('=')[1].strip().strip('"\'') break if external_tool_name: # Find the class name from the file for line in content.split('\n'): if 'class ' in line and 'BaseTool' in line: class_name = line.split('class ')[1].split('(')[0].strip() # Build both mappings short_to_long[class_name] = external_tool_name long_to_internal[external_tool_name] = { "class_name": class_name, "dir_name": dir_name } print(f"Mapped: {class_name} -> {external_tool_name} (dir: {dir_name})") break except Exception as e: print(f"Warning: Could not extract TOOL_NAME from {tool_file_path}: {str(e)}") continue return {"short_to_long": short_to_long, "long_to_internal": long_to_internal} def _load_single_tool(self, root: str, import_path: str, agentflow_dir: str) -> Dict[str, Any]: """ Load all tools from a single module and return their metadata. This method is designed to be called in parallel. Note: A single module may contain multiple tool classes (either defined locally or imported). Returns: Dict with lists of tool metadata/instances, or error information """ result = {'metadata_list': [], 'instance_list': [], 'errors': []} try: module = importlib.import_module(import_path) current_dir_name = os.path.basename(root) for name, obj in inspect.getmembers(module): if inspect.isclass(obj) and name.endswith('Tool') and name != 'BaseTool': try: # Check if the tool requires specific llm engine tool_index = -1 for i, tool_name in enumerate(self.enabled_tools): # First check short_to_long mapping if hasattr(self, 'tool_name_mapping'): short_to_long = self.tool_name_mapping.get('short_to_long', {}) long_to_internal = self.tool_name_mapping.get('long_to_internal', {}) # If input is short name, convert to long name long_name = short_to_long.get(tool_name, tool_name) # Check if long name matches this directory if long_name in long_to_internal: if long_to_internal[long_name]["dir_name"] == current_dir_name: tool_index = i break # Fallback to original behavior if tool_name.lower().replace('_tool', '') == current_dir_name: tool_index = i break if tool_index >= 0 and tool_index < len(self.tool_engine): engine = self.tool_engine[tool_index] if engine == "Default": tool_instance = obj() elif engine == "self": tool_instance = obj(model_string=self.model_string) else: tool_instance = obj(model_string=engine) else: tool_instance = obj() # Use the external tool name (from TOOL_NAME) as the key metadata_key = getattr(tool_instance, 'tool_name', name) metadata = { 'tool_name': getattr(tool_instance, 'tool_name', 'Unknown'), 'tool_description': getattr(tool_instance, 'tool_description', 'No description'), 'tool_version': getattr(tool_instance, 'tool_version', 'Unknown'), 'input_types': getattr(tool_instance, 'input_types', {}), 'output_type': getattr(tool_instance, 'output_type', 'Unknown'), 'demo_commands': getattr(tool_instance, 'demo_commands', []), 'user_metadata': getattr(tool_instance, 'user_metadata', {}), 'require_llm_engine': getattr(obj, 'require_llm_engine', False), } result['metadata_list'].append((metadata_key, metadata)) result['instance_list'].append((metadata_key, tool_instance)) except Exception as e: result['errors'].append(f"Error instantiating {name}: {str(e)}") except Exception as e: result['errors'].append(f"Error loading module {import_path}: {str(e)}") return result def load_tools_and_get_metadata(self, parallel: bool = True, max_workers: int = 4) -> Dict[str, Any]: """ Load tools and get metadata. Can be done in parallel for faster initialization. Args: parallel: If True, load tools in parallel using ThreadPoolExecutor max_workers: Maximum number of worker threads (default: 4) """ print(f"Loading tools and getting metadata... (parallel={parallel}, max_workers={max_workers})") start_time = time.time() self.toolbox_metadata = {} agentflow_dir = self.get_project_root() tools_dir = os.path.join(agentflow_dir, 'tools') # Add the agentflow directory and its parent to the Python path sys.path.insert(0, agentflow_dir) sys.path.insert(0, os.path.dirname(agentflow_dir)) print(f"Updated Python path: {sys.path}") if not os.path.exists(tools_dir): print(f"Error: Tools directory does not exist: {tools_dir}") return self.toolbox_metadata # Build tool name mapping if not already built if not hasattr(self, 'tool_name_mapping'): self.tool_name_mapping = self.build_tool_name_mapping(tools_dir) print(f"\n==> Tool name mapping (short to long): {self.tool_name_mapping.get('short_to_long', {})}") print(f"==> Tool name mapping (long to internal): {self.tool_name_mapping.get('long_to_internal', {})}") # Collect all tool directories to process, maintaining the order from available_tools tool_dirs_to_process = [] if self.load_all: # If loading all tools, use os.walk order for root, dirs, files in os.walk(tools_dir): if 'tool.py' in files: file = 'tool.py' module_path = os.path.join(root, file) relative_path = os.path.relpath(module_path, agentflow_dir) import_path = '.'.join(os.path.split(relative_path)).replace(os.sep, '.')[:-3] tool_dirs_to_process.append((root, import_path)) else: # Build a map of directory names to paths for efficient lookup dir_to_paths = {} for root, dirs, files in os.walk(tools_dir): if 'tool.py' in files: dir_name = os.path.basename(root) file = 'tool.py' module_path = os.path.join(root, file) relative_path = os.path.relpath(module_path, agentflow_dir) import_path = '.'.join(os.path.split(relative_path)).replace(os.sep, '.')[:-3] dir_to_paths[dir_name] = (root, import_path) # Process in the order of available_tools (which matches enabled_tools order) for tool_dir in self.available_tools: if tool_dir in dir_to_paths: tool_dirs_to_process.append(dir_to_paths[tool_dir]) else: print(f"Warning: Tool directory '{tool_dir}' not found in tools directory") if parallel and len(tool_dirs_to_process) > 1: # Parallel loading print(f"\n==> Loading {len(tool_dirs_to_process)} tool modules in parallel...") with ThreadPoolExecutor(max_workers=max_workers) as executor: # Submit all tool loading tasks and maintain order futures = [ (executor.submit(self._load_single_tool, root, import_path, agentflow_dir), import_path) for root, import_path in tool_dirs_to_process ] # Wait for all to complete, then process in original order print(f"Waiting for all {len(futures)} modules to load...") # Process results in the original submission order for future, import_path in futures: try: result = future.result() # Report any errors if result['errors']: for error in result['errors']: print(f"Error loading {import_path}: {error}") # Process all tools found in this module (in the order they appear in the module) for metadata_key, metadata in result['metadata_list']: self.toolbox_metadata[metadata_key] = metadata for instance_key, instance in result['instance_list']: self.tool_instances_cache[instance_key] = instance print(f"Loaded: {instance_key} with engine: {getattr(instance, 'model_string', 'default')}") except Exception as e: print(f"Exception loading {import_path}: {str(e)}") else: # Serial loading (original behavior) print(f"\n==> Loading {len(tool_dirs_to_process)} tool modules serially...") for root, import_path in tool_dirs_to_process: print(f"\n==> Attempting to import: {import_path}") result = self._load_single_tool(root, import_path, agentflow_dir) # Report any errors if result['errors']: for error in result['errors']: print(f"Error: {error}") # Process all tools found in this module for metadata_key, metadata in result['metadata_list']: self.toolbox_metadata[metadata_key] = metadata print(f"Metadata for {metadata_key}: {metadata}") for instance_key, instance in result['instance_list']: self.tool_instances_cache[instance_key] = instance print(f"Cached tool instance: {instance_key} with engine: {getattr(instance, 'model_string', 'default')}") elapsed_time = time.time() - start_time print(f"\n==> Total number of tools imported: {len(self.toolbox_metadata)} (took {elapsed_time:.2f}s)") return self.toolbox_metadata def run_demo_commands(self) -> List[str]: print("\n==> Running demo commands for each tool...") self.available_tools = [] # Process tools in alphabetical order by tool name for tool_name, tool_data in sorted(self.toolbox_metadata.items()): print(f"Checking availability of {tool_name}...") try: # Use the cached tool instance instead of creating a new one # This preserves the engine configuration from load_tools_and_get_metadata if tool_name in self.tool_instances_cache: tool_instance = self.tool_instances_cache[tool_name] print(f"Using cached instance with engine: {getattr(tool_instance, 'model_string', 'default')}") else: # Fallback: create new instance if not in cache # tool_name here is the long external name from metadata # We need to get the internal class name and directory if hasattr(self, 'tool_name_mapping'): long_to_internal = self.tool_name_mapping.get('long_to_internal', {}) if tool_name in long_to_internal: dir_name = long_to_internal[tool_name]["dir_name"] class_name = long_to_internal[tool_name]["class_name"] else: # Fallback to original behavior dir_name = tool_name.lower().replace('_tool', '') class_name = tool_name else: # Fallback to original behavior dir_name = tool_name.lower().replace('_tool', '') class_name = tool_name # Import the tool module module_name = f"tools.{dir_name}.tool" module = importlib.import_module(module_name) # Get the tool class tool_class = getattr(module, class_name) # Instantiate the tool tool_instance = tool_class() print(f"Created new instance (not in cache)") # FIXME This is a temporary workaround to avoid running demo commands self.available_tools.append(tool_name) except Exception as e: print(f"Error checking availability of {tool_name}: {str(e)}") print(traceback.format_exc()) # update the toolmetadata with the available tools self.toolbox_metadata = {tool: self.toolbox_metadata[tool] for tool in self.available_tools} print("\nFinished running demo commands for each tool.") # print(f"Updated total number of available tools: {len(self.toolbox_metadata)}") # print(f"Available tools: {self.available_tools}") return self.available_tools def _set_up_tools(self) -> None: print("\n==> Setting up tools...") # First, build a temporary mapping by scanning all tools agentflow_dir = self.get_project_root() tools_dir = os.path.join(agentflow_dir, 'tools') self.tool_name_mapping = self.build_tool_name_mapping(tools_dir) if os.path.exists(tools_dir) else {} # Map input tool names (short) to internal directory names for filtering mapped_tools = [] short_to_long = self.tool_name_mapping.get('short_to_long', {}) long_to_internal = self.tool_name_mapping.get('long_to_internal', {}) for i, tool in enumerate(self.enabled_tools): # If tool is a short name, convert to long name first long_name = short_to_long.get(tool, tool) print(f" [{i}] {tool} -> {long_name}", end="") # Then get the directory name if long_name in long_to_internal: dir_name = long_to_internal[long_name]["dir_name"] mapped_tools.append(dir_name) print(f" -> {dir_name}") else: # Fallback to original behavior for unmapped tools dir_name = tool.lower().replace('_tool', '') mapped_tools.append(dir_name) print(f" -> {dir_name} (fallback)") self.available_tools = mapped_tools print(f"\n==> Mapped tools (directory names): {mapped_tools}") # Now load tools and get metadata (with optional parallel loading) self.load_tools_and_get_metadata( parallel=self.parallel_loading, max_workers=self.max_workers ) # Run demo commands to determine available tools # This will update self.available_tools to contain external names self.run_demo_commands() # available_tools is now already updated by run_demo_commands with external names print("Finished setting up tools.") print(f"Total number of final available tools: {len(self.available_tools)}") print(f"Final available tools: {self.available_tools}") if __name__ == "__main__": import time enabled_tools = ["Base_Generator_Tool", "Python_Coder_Tool", "SearXNG_Search_Tool"] tool_engine = ["gpt-4o-mini", "gpt-4o-mini", "Default", "Default"] print("\n" + "="*80) print("PERFORMANCE COMPARISON: Serial vs Parallel Tool Loading") print("="*80) # Test 1: Serial loading print("\n[Test 1] Serial Loading...") print("-"*80) start = time.time() init_serial = Initializer( enabled_tools=enabled_tools, tool_engine=tool_engine, parallel_loading=False # Serial ) serial_time = time.time() - start print(f"\nSerial loading completed in {serial_time:.2f}s") print(f" Tools loaded: {len(init_serial.available_tools)}") # Test 2: Parallel loading (4 workers) print("\n[Test 2] Parallel Loading (4 workers)...") print("-"*80) start = time.time() init_parallel = Initializer( enabled_tools=enabled_tools, tool_engine=tool_engine, parallel_loading=True, # Parallel max_workers=4 ) parallel_time = time.time() - start print(f"\nParallel loading completed in {parallel_time:.2f}s") print(f" Tools loaded: {len(init_parallel.available_tools)}") # Summary print("\n" + "="*80) print("PERFORMANCE SUMMARY") print("="*80) print(f"Serial loading: {serial_time:>6.2f}s") print(f"Parallel loading: {parallel_time:>6.2f}s (4 workers)") print(f"Speedup: {serial_time/parallel_time:>6.2f}x") print(f"Time saved: {serial_time - parallel_time:>6.2f}s") print("="*80)