PraisonAI 3.0.0__py3-none-any.whl

praisonai/templates/tool_override.py

@@ -0,0 +1,498 @@
+"""
+Tool Override Loader for PraisonAI.
+
+Allows loading custom tools from files, modules, and directories.
+Supports runtime tool registration with context manager pattern.
+"""
+
+import ast
+import importlib.util
+import sys
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, Callable, Dict, Generator, List, Optional
+
+
+class SecurityError(Exception):
+    """Raised when a security violation is detected."""
+    pass
+
+
+class ToolOverrideLoader:
+    """
+    Loads custom tools from various sources.
+    
+    Supports:
+    - Python file paths
+    - Module import paths
+    - Directories containing tool files
+    
+    Security:
+    - Only local paths allowed by default
+    - Remote URLs rejected
+    - No auto-execution of arbitrary code
+    """
+    
+    # Default custom tool directories
+    DEFAULT_TOOL_DIRS = [
+        "~/.praison/tools",
+        "~/.config/praison/tools",
+    ]
+    
+    def __init__(self):
+        """Initialize tool override loader."""
+        self._loaded_tools: Dict[str, Callable] = {}
+    
+    def get_default_tool_dirs(self) -> List[Path]:
+        """
+        Get default custom tool directories.
+        
+        Returns:
+            List of Path objects for default tool directories
+        """
+        return [Path(d).expanduser() for d in self.DEFAULT_TOOL_DIRS]
+    
+    def load_from_file(self, file_path: str) -> Dict[str, Callable]:
+        """
+        Load tools from a Python file.
+        
+        Args:
+            file_path: Path to Python file containing tool functions
+            
+        Returns:
+            Dict mapping tool names to callable functions
+            
+        Raises:
+            SecurityError: If path is a remote URL
+            FileNotFoundError: If file doesn't exist
+        """
+        # Security check - reject remote URLs
+        if file_path.startswith(("http://", "https://", "ftp://")):
+            raise SecurityError(
+                f"Remote URLs not allowed for security: {file_path}. "
+                "Only local file paths are permitted."
+            )
+        
+        path = Path(file_path).expanduser().resolve()
+        
+        if not path.exists():
+            raise FileNotFoundError(f"Tool file not found: {path}")
+        
+        if not path.suffix == ".py":
+            raise ValueError(f"Tool file must be a Python file (.py): {path}")
+        
+        # Load the module
+        spec = importlib.util.spec_from_file_location(
+            f"custom_tools_{path.stem}",
+            path
+        )
+        if spec is None or spec.loader is None:
+            raise ImportError(f"Could not load module from: {path}")
+        
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[spec.name] = module
+        spec.loader.exec_module(module)
+        
+        # Extract callable functions (tools)
+        tools = {}
+        for name in dir(module):
+            if name.startswith("_"):
+                continue
+            obj = getattr(module, name)
+            if callable(obj) and not isinstance(obj, type):
+                tools[name] = obj
+        
+        self._loaded_tools.update(tools)
+        return tools
+    
+    def load_from_module(self, module_path: str) -> Dict[str, Callable]:
+        """
+        Load tools from a module import path.
+        
+        Args:
+            module_path: Python module path (e.g., 'mypackage.tools')
+            
+        Returns:
+            Dict mapping tool names to callable functions
+        """
+        try:
+            module = importlib.import_module(module_path)
+        except ImportError as e:
+            raise ImportError(f"Could not import module: {module_path}") from e
+        
+        # Extract callable functions (tools)
+        tools = {}
+        for name in dir(module):
+            if name.startswith("_"):
+                continue
+            obj = getattr(module, name)
+            if callable(obj) and not isinstance(obj, type):
+                tools[name] = obj
+        
+        self._loaded_tools.update(tools)
+        return tools
+    
+    def load_from_directory(self, dir_path: str) -> Dict[str, Callable]:
+        """
+        Load tools from all Python files in a directory.
+        
+        Args:
+            dir_path: Path to directory containing tool files
+            
+        Returns:
+            Dict mapping tool names to callable functions
+            
+        Raises:
+            SecurityError: If path is a remote URL
+        """
+        # Security check
+        if dir_path.startswith(("http://", "https://", "ftp://")):
+            raise SecurityError(
+                f"Remote URLs not allowed for security: {dir_path}. "
+                "Only local directory paths are permitted."
+            )
+        
+        path = Path(dir_path).expanduser().resolve()
+        
+        if not path.exists():
+            raise FileNotFoundError(f"Tool directory not found: {path}")
+        
+        if not path.is_dir():
+            raise ValueError(f"Path is not a directory: {path}")
+        
+        tools = {}
+        for py_file in path.glob("*.py"):
+            if py_file.name.startswith("_"):
+                continue
+            try:
+                file_tools = self.load_from_file(str(py_file))
+                tools.update(file_tools)
+            except Exception:
+                # Skip files that fail to load
+                pass
+        
+        return tools
+    
+    def discover_tools_in_directory(self, dir_path: str) -> List[str]:
+        """
+        Discover tool names in a directory without executing code.
+        
+        Uses AST parsing to find function definitions.
+        
+        Args:
+            dir_path: Path to directory to scan
+            
+        Returns:
+            List of discovered tool function names
+        """
+        path = Path(dir_path).expanduser().resolve()
+        
+        if not path.exists() or not path.is_dir():
+            return []
+        
+        tool_names = []
+        for py_file in path.glob("*.py"):
+            if py_file.name.startswith("_"):
+                continue
+            try:
+                content = py_file.read_text()
+                tree = ast.parse(content)
+                
+                for node in ast.walk(tree):
+                    if isinstance(node, ast.FunctionDef):
+                        if not node.name.startswith("_"):
+                            tool_names.append(node.name)
+            except Exception:
+                # Skip files that fail to parse
+                pass
+        
+        return tool_names
+    
+    @contextmanager
+    def override_context(
+        self,
+        files: Optional[List[str]] = None,
+        modules: Optional[List[str]] = None,
+        directories: Optional[List[str]] = None
+    ) -> Generator[Dict[str, Callable], None, None]:
+        """
+        Context manager for temporary tool overrides.
+        
+        Tools loaded within this context are available only during
+        the context and are cleaned up afterward.
+        
+        Args:
+            files: List of Python file paths to load
+            modules: List of module paths to import
+            directories: List of directories to scan
+            
+        Yields:
+            Dict of loaded tools
+        """
+        # Store original state
+        original_tools = self._loaded_tools.copy()
+        
+        try:
+            # Load tools from all sources
+            loaded = {}
+            
+            if files:
+                for f in files:
+                    loaded.update(self.load_from_file(f))
+            
+            if modules:
+                for m in modules:
+                    loaded.update(self.load_from_module(m))
+            
+            if directories:
+                for d in directories:
+                    loaded.update(self.load_from_directory(d))
+            
+            yield loaded
+            
+        finally:
+            # Restore original state
+            self._loaded_tools = original_tools
+    
+    def get_loaded_tools(self) -> Dict[str, Callable]:
+        """
+        Get all currently loaded tools.
+        
+        Returns:
+            Dict mapping tool names to callable functions
+        """
+        return self._loaded_tools.copy()
+    
+    def clear_loaded_tools(self) -> None:
+        """Clear all loaded tools."""
+        self._loaded_tools.clear()
+
+
+def create_tool_registry_with_overrides(
+    override_files: Optional[List[str]] = None,
+    override_dirs: Optional[List[str]] = None,
+    include_defaults: bool = True,
+    tools_sources: Optional[List[str]] = None,
+    template_dir: Optional[str] = None,
+) -> Dict[str, Callable]:
+    """
+    Create a tool registry with custom overrides.
+    
+    Resolution order (highest priority first):
+    1. Override files (explicit CLI --tools)
+    2. Override directories (explicit CLI --tools-dir)
+    3. Template tools_sources (from TEMPLATE.yaml)
+    4. Template-local tools.py
+    4.5. Current working directory tools.py (./tools.py)
+    5. Default custom dirs (~/.praison/tools, etc.)
+    6. Package discovery (praisonai-tools if installed)
+    7. Built-in tools
+    
+    Args:
+        override_files: Explicit tool files to load
+        override_dirs: Directories to scan for tools
+        include_defaults: Whether to include default tool directories
+        tools_sources: Template-declared tool sources (modules or paths)
+        template_dir: Template directory for local tools.py
+        
+    Returns:
+        Dict mapping tool names to callable functions
+    """
+    registry = {}
+    loader = ToolOverrideLoader()
+    
+    # 7. Start with built-in tools (lowest priority)
+    try:
+        from praisonaiagents.tools import TOOL_MAPPINGS
+        registry.update(TOOL_MAPPINGS)
+    except ImportError:
+        pass
+    
+    # 6. Package discovery - try praisonai-tools if installed
+    try:
+        import praisonai_tools.tools as external_tools
+        # Get all exported tools from praisonai_tools
+        for name in dir(external_tools):
+            if not name.startswith('_'):
+                obj = getattr(external_tools, name, None)
+                if callable(obj) or (hasattr(obj, 'run') and callable(getattr(obj, 'run', None))):
+                    registry[name] = obj
+    except ImportError:
+        pass
+    
+    # 5. Add default custom dirs
+    if include_defaults:
+        for dir_path in loader.get_default_tool_dirs():
+            if dir_path.exists():
+                try:
+                    tools = loader.load_from_directory(str(dir_path))
+                    registry.update(tools)
+                except Exception:
+                    pass
+    
+    # 4.5. Current working directory tools.py (if exists)
+    cwd_tools_py = Path.cwd() / "tools.py"
+    if cwd_tools_py.exists():
+        try:
+            tools = loader.load_from_file(str(cwd_tools_py))
+            registry.update(tools)
+        except Exception:
+            pass
+    
+    # 4. Template-local tools.py
+    if template_dir:
+        tools_py = Path(template_dir) / "tools.py"
+        if tools_py.exists():
+            try:
+                tools = loader.load_from_file(str(tools_py))
+                registry.update(tools)
+            except Exception:
+                pass
+    
+    # 3. Template tools_sources (from TEMPLATE.yaml)
+    if tools_sources:
+        for source in tools_sources:
+            try:
+                # Security: only allow local paths and python modules
+                if source.startswith(("http://", "https://", "ftp://")):
+                    continue  # Skip remote URLs
+                
+                source_path = Path(source).expanduser()
+                
+                if source_path.exists():
+                    # It's a local path
+                    if source_path.is_file() and source_path.suffix == ".py":
+                        tools = loader.load_from_file(str(source_path))
+                        registry.update(tools)
+                    elif source_path.is_dir():
+                        tools = loader.load_from_directory(str(source_path))
+                        registry.update(tools)
+                else:
+                    # Try as a Python module path
+                    tools = loader.load_from_module(source)
+                    registry.update(tools)
+            except Exception:
+                pass
+    
+    # 2. Add override directories (CLI --tools-dir)
+    if override_dirs:
+        for dir_path in override_dirs:
+            try:
+                tools = loader.load_from_directory(dir_path)
+                registry.update(tools)
+            except Exception:
+                pass
+    
+    # 1. Add override files (highest priority, CLI --tools)
+    if override_files:
+        for file_path in override_files:
+            try:
+                tools = loader.load_from_file(file_path)
+                registry.update(tools)
+            except Exception:
+                pass
+    
+    return registry
+
+
+def resolve_tools(
+    tool_names: List[Any],
+    registry: Optional[Dict[str, Callable]] = None,
+    template_dir: Optional[str] = None,
+) -> List[Callable]:
+    """
+    Resolve tool names to callable tools from registry.
+    
+    Handles:
+    - String tool names (looked up in registry)
+    - Already-callable tools (passed through)
+    - Built-in tool names (shell_tool, file_tool, etc.)
+    
+    Args:
+        tool_names: List of tool names (strings) or callables
+        registry: Tool registry to look up names in
+        template_dir: Optional template directory for local tools.py autoload
+        
+    Returns:
+        List of resolved callable tools
+    """
+    if not tool_names:
+        return []
+    
+    resolved = []
+    
+    # Build registry if not provided
+    if registry is None:
+        registry = create_tool_registry_with_overrides(include_defaults=True)
+    
+    # Load template-local tools.py if exists
+    if template_dir:
+        loader = ToolOverrideLoader()
+        tools_py = Path(template_dir) / "tools.py"
+        if tools_py.exists():
+            try:
+                local_tools = loader.load_from_file(str(tools_py))
+                registry.update(local_tools)
+            except Exception:
+                pass
+    
+    for tool in tool_names:
+        if callable(tool):
+            # Already a callable, use directly
+            resolved.append(tool)
+        elif isinstance(tool, str):
+            # Look up by name in registry
+            tool_name = tool.strip()
+            
+            # Try exact match first
+            if tool_name in registry:
+                tool_obj = registry[tool_name]
+                # Handle lazy-loaded tools (tuples of module, class)
+                if isinstance(tool_obj, tuple):
+                    try:
+                        module_name, class_name = tool_obj
+                        import importlib
+                        module = importlib.import_module(module_name)
+                        tool_class = getattr(module, class_name)
+                        resolved.append(tool_class())
+                    except Exception:
+                        pass
+                elif callable(tool_obj):
+                    resolved.append(tool_obj)
+                else:
+                    # Try to instantiate if it's a class
+                    try:
+                        resolved.append(tool_obj())
+                    except Exception:
+                        resolved.append(tool_obj)
+            else:
+                # Try common variations
+                variations = [
+                    tool_name,
+                    tool_name.lower(),
+                    tool_name.replace("-", "_"),
+                    tool_name.replace("_", "-"),
+                    f"{tool_name}_tool",
+                    f"{tool_name}Tool",
+                ]
+                found = False
+                for var in variations:
+                    if var in registry:
+                        tool_obj = registry[var]
+                        if callable(tool_obj):
+                            resolved.append(tool_obj)
+                            found = True
+                            break
+                
+                if not found:
+                    # Try to import from praisonaiagents.tools
+                    try:
+                        from praisonaiagents import tools as agent_tools
+                        if hasattr(agent_tools, tool_name):
+                            resolved.append(getattr(agent_tools, tool_name))
+                        elif hasattr(agent_tools, f"{tool_name}_tool"):
+                            resolved.append(getattr(agent_tools, f"{tool_name}_tool"))
+                    except (ImportError, AttributeError):
+                        pass
+    
+    return resolved