tinyagent-py 0.0.13__py3-none-any.whl → 0.0.16__py3-none-any.whl

tinyagent/code_agent/helper.py

@@ -47,13 +47,13 @@ You are an Agent, You need to solve the task, not suggesting user about how to s
 
                             """)
 
-def load_template(path: str) -> str:
+def load_template(path: str,key:str="system_prompt") -> str:
     """
     Load the YAML file and extract its 'system_prompt' field.
     """
     with open(path, "r") as f:
         data = yaml.safe_load(f)
-    return data["system_prompt"]
+    return data[key]
 
 def render_system_prompt(template_str: str,
                          tools: dict,

tinyagent/code_agent/modal_sandbox.py

@@ -78,7 +78,7 @@ def create_sandbox(
 
     if apt_packages is None:
         # Always install the basics required for most workflows
-        apt_packages = ("git", "curl", "nodejs", "npm")
+        apt_packages = ("git", "curl", "nodejs", "npm","ripgrep","tree")
 
     if default_packages is None:
         default_packages = (

tinyagent/code_agent/providers/__init__.py

@@ -1,4 +1,17 @@
 from .base import CodeExecutionProvider
 from .modal_provider import ModalProvider
 
-__all__ = ["CodeExecutionProvider", "ModalProvider"] 
+# Import SeatbeltProvider conditionally to avoid errors on non-macOS systems
+import platform
+if platform.system() == "Darwin":
+    try:
+        from .seatbelt_provider import SeatbeltProvider
+    except ImportError:
+        # If there's an issue importing, just don't make it available
+        pass
+
+__all__ = ["CodeExecutionProvider", "ModalProvider"]
+
+# Add SeatbeltProvider to __all__ if it was successfully imported
+if platform.system() == "Darwin" and "SeatbeltProvider" in globals():
+    __all__.append("SeatbeltProvider") 

tinyagent/code_agent/providers/base.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Dict, List, Any, Optional
+from typing import Dict, List, Any, Optional, Set
 from tinyagent.hooks.logging_manager import LoggingManager
 import cloudpickle
 
@@ -21,6 +21,9 @@ class CodeExecutionProvider(ABC):
         pip_packages: List[str] = None,
         secrets: Dict[str, Any] = None,
         lazy_init: bool = True,
+        bypass_shell_safety: bool = False,
+        additional_safe_shell_commands: Optional[List[str]] = None,
+        additional_safe_control_operators: Optional[List[str]] = None,
         **kwargs
     ):
         self.log_manager = log_manager
@@ -35,6 +38,36 @@ class CodeExecutionProvider(ABC):
         self._locals_dict = kwargs.get("locals_dict", {})
         self._user_variables = {}
         self.code_tools_definitions = []
+        
+        # Shell safety configuration
+        self.bypass_shell_safety = bypass_shell_safety
+        
+        # Safe shell commands that don't modify the system or access sensitive data
+        self.safe_shell_commands: Set[str] = {
+            "ls", "cat", "grep", "find", "echo", "pwd", "whoami", "date", 
+            "head", "tail", "wc", "sort", "uniq", "tr", "cut", "sed", "awk",
+            "ps", "df", "du", "uname", "which", "type", "file", "stat", "rg", "if",
+            "tree"
+        }
+        
+        # Add additional safe shell commands if provided
+        if additional_safe_shell_commands:
+            if "*" in additional_safe_shell_commands:
+                # If wildcard is provided, allow all commands (effectively bypassing the check)
+                self.bypass_shell_safety = True
+            else:
+                self.safe_shell_commands.update(additional_safe_shell_commands)
+        
+        # Safe control operators for shell commands
+        self.safe_control_operators: Set[str] = {"&&", "||", ";", "|"}
+        
+        # Add additional safe control operators if provided
+        if additional_safe_control_operators:
+            if "*" in additional_safe_control_operators:
+                # If wildcard is provided, allow all operators
+                self.safe_control_operators = set("*")
+            else:
+                self.safe_control_operators.update(additional_safe_control_operators)
     
     @abstractmethod
     async def execute_python(
@@ -58,6 +91,133 @@ class CodeExecutionProvider(ABC):
         """
         pass
     
+    @abstractmethod
+    async def execute_shell(
+        self,
+        command: List[str],
+        timeout: int = 10,
+        workdir: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Execute a shell command securely and return the result.
+        
+        Args:
+            command: List of command parts to execute
+            timeout: Maximum execution time in seconds
+            workdir: Working directory for command execution
+            
+        Returns:
+            Dictionary containing execution results with keys:
+            - stdout: stdout from the execution
+            - stderr: stderr from the execution
+            - exit_code: exit code from the command
+        """
+        pass
+    
+    def is_safe_command(self, command: List[str]) -> Dict[str, Any]:
+        """
+        Check if a shell command is safe to execute.
+        
+        Args:
+            command: List of command parts to check
+            
+        Returns:
+            Dictionary with:
+            - safe: Boolean indicating if command is safe
+            - reason: Reason why command is not safe (if applicable)
+        """
+        # If shell safety checks are bypassed, consider all commands safe
+        if self.bypass_shell_safety:
+            return {"safe": True}
+            
+        if type(command) == str:
+            command = command.split(" ")
+        if not command or not isinstance(command, list) or len(command) == 0:
+            return {"safe": False, "reason": "Empty or invalid command"}
+        
+        # Special handling for bash -c or bash -lc commands
+        if len(command) >= 3 and command[0] == "bash" and command[1] in ["-c", "-lc"]:
+            # For bash -c or bash -lc, we need to parse the command string that follows
+            # We'll extract commands from the bash command string and check them
+            bash_cmd_str = command[2]
+            
+            # Simple parsing of the bash command to extract command names
+            # This is a basic implementation and might not cover all edge cases
+            import shlex
+            import re
+            
+            try:
+                # Shell script keywords that should be allowed
+                shell_keywords = {
+                    "if", "then", "else", "elif", "fi", "for", "do", "done", 
+                    "while", "until", "case", "esac", "in", "function", "select",
+                    "time", "coproc", "true", "false"
+                }
+                
+                # Split the command by common shell operators
+                cmd_parts = re.split(r'(\||;|&&|\|\||>|>>|<|<<)', bash_cmd_str)
+                commands_to_check = []
+                
+                for part in cmd_parts:
+                    part = part.strip()
+                    if part and part not in ['|', ';', '&&', '||', '>', '>>', '<', '<<']:
+                        # Get the first word which is typically the command
+                        try:
+                            words = shlex.split(part)
+                            if words:
+                                cmd_name = words[0].split('/')[-1]  # Extract binary name
+                                
+                                # Skip shell keywords
+                                if cmd_name in shell_keywords:
+                                    continue
+                                    
+                                # Skip variable assignments (e.g., VAR=value)
+                                if re.match(r'^[A-Za-z_][A-Za-z0-9_]*=', cmd_name):
+                                    continue
+                                
+                                if cmd_name not in self.safe_shell_commands and '*' not in cmd_name and '?' not in cmd_name:
+                                    return {"safe": False, "reason": f"Unsafe command in bash script: {cmd_name}"}
+                        except Exception:
+                            # If parsing fails, be cautious and reject
+                            return {"safe": False, "reason": "Could not parse bash command safely"}
+                
+                # All commands in the bash script are safe
+                return {"safe": True}
+            except Exception as e:
+                return {"safe": False, "reason": f"Error parsing bash command: {str(e)}"}
+        
+        # Normal command processing for non-bash -c commands
+        # Shell operators that might be passed as separate arguments
+        shell_operators = ['|', '>', '<', '>>', '<<', '&&', '||', ';']
+        
+        # Extract actual commands from the command list, ignoring shell operators
+        commands_to_check = []
+        i = 0
+        while i < len(command):
+            if command[i] in shell_operators:
+                i += 1
+                continue
+            
+            # Extract the binary name
+            bin_name = command[i].split("/")[-1]
+            commands_to_check.append(bin_name)
+            
+            # Skip to next command after an operator
+            i += 1
+            while i < len(command) and command[i] not in shell_operators:
+                i += 1
+        
+        # Check if all commands are in the safe list
+        for cmd in commands_to_check:
+            # Handle wildcards in command names (e.g., *.py)
+            if '*' in cmd or '?' in cmd:
+                continue
+                
+            if cmd not in self.safe_shell_commands:
+                return {"safe": False, "reason": f"Unsafe command: {cmd}"}
+        
+        return {"safe": True}
+    
     @abstractmethod
     async def cleanup(self):
         """Clean up any resources used by the provider."""
@@ -129,14 +289,14 @@ class CodeExecutionProvider(ABC):
         if variables_str_list:
             # Find where to insert (after tools section if it exists)
             insert_index = 0
-            for i, code in enumerate(self.default_python_codes):
+            for i, code in enumerate(self.code_tools_definitions):
                 if "###########</tools>###########" in code:
                     insert_index = i + 1
                     break
             
             # Insert the variables code
             for j, var_code in enumerate(variables_str_list):
-                self.default_python_codes.insert(insert_index + j, var_code)
+                self.code_tools_definitions.insert(insert_index + j, var_code)
     
     def _remove_existing_user_variables(self) -> None:
         """Remove existing user variables from default python codes."""
@@ -144,16 +304,16 @@ class CodeExecutionProvider(ABC):
         start_index = None
         end_index = None
         
-        for i, code in enumerate(self.default_python_codes):
+        for i, code in enumerate(self.code_tools_definitions):
             if "###########<user_variables>###########" in code:
-                start_index = i - 1 if i > 0 and "import cloudpickle" in self.default_python_codes[i-1] else i
+                start_index = i - 1 if i > 0 and "import cloudpickle" in self.code_tools_definitions[i-1] else i
             elif "###########</user_variables>###########" in code:
                 end_index = i + 2  # Include the newline after
                 break
         
         if start_index is not None and end_index is not None:
             # Remove the old variables section
-            del self.default_python_codes[start_index:end_index]
+            del self.code_tools_definitions[start_index:end_index]
     
     def get_user_variables(self) -> Dict[str, Any]:
         """
@@ -204,4 +364,18 @@ class CodeExecutionProvider(ABC):
                 self._user_variables[var_name] = var_value
             except Exception:
                 # If serialization fails, skip this variable
-                pass 
+                pass 
+    
+    def shell_response_to_llm_understandable(self, response: Dict[str, Any]) -> str:
+        """
+        Convert a shell command response to a format that is understandable by the LLM.
+        """
+        if response.get('stderr',None) not in [None,""]:
+            error_message = "Bash Error: " + response['stderr']
+            if "No such file or directory" in response['stderr']:
+                error_message.replace("No such file or directory", "No such file or directory, Have you provided the correct absolute path? If you are unsure use ls first to make sure the path exists")
+            if "Command timed out after" in response['stderr']:
+                error_message += ", Make sure your command is specific enough. And only if it is the most specific and optimized command then try to increase the timeout parameter if you need to more time for this command."
+            return error_message
+        else:
+            return response['stdout']

tinyagent/code_agent/providers/modal_provider.py

@@ -1,9 +1,22 @@
 import sys
 import modal
 import cloudpickle
+from pprint import pprint
 from typing import Dict, List, Any, Optional, Union
 from .base import CodeExecutionProvider
-from ..utils import clean_response, make_session_blob, _run_python
+from ..utils import clean_response, make_session_blob, _run_python, _run_shell
+try:
+    from ..modal_sandbox import COLOR
+except ImportError:
+    # Fallback colors if modal_sandbox is not available
+    COLOR = {
+    "HEADER": "\033[95m",
+    "BLUE": "\033[94m",
+    "GREEN": "\033[92m",
+    "RED": "\033[91m",
+    "ENDC": "\033[0m",
+}
+
 
 
 class ModalProvider(CodeExecutionProvider):
@@ -16,6 +29,7 @@ class ModalProvider(CodeExecutionProvider):
     """
     
     PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}"
+    TIMEOUT_MAX = 120
     
     def __init__(
         self,
@@ -27,25 +41,45 @@ class ModalProvider(CodeExecutionProvider):
         apt_packages: Optional[List[str]] = None,
         python_version: Optional[str] = None,
         authorized_imports: list[str] | None = None,
+        authorized_functions: list[str] | None = None,
         modal_secrets: Dict[str, Union[str, None]] | None = None,
         lazy_init: bool = True,
         sandbox_name: str = "tinycodeagent-sandbox",
         local_execution: bool = False,
+        check_string_obfuscation: bool = True,
+        bypass_shell_safety: bool = False,  # Default to False for ModalProvider
+        additional_safe_shell_commands: Optional[List[str]] = None,
+        additional_safe_control_operators: Optional[List[str]] = None,
         **kwargs
     ):
-        """Create a ModalProvider instance.
-
-        Additional keyword arguments (passed via **kwargs) are ignored by the
-        base class but accepted here for forward-compatibility.
-
+        """
+        Initialize Modal-based code execution provider.
+        
         Args:
-            default_packages: Base set of Python packages installed into the
-                sandbox image. If ``None`` a sane default list is used. The
-                final set of installed packages is the union of
-                ``default_packages`` and ``pip_packages``.
-            apt_packages: Debian/Ubuntu APT packages to install into the image
-                prior to ``pip install``. Defaults to an empty list.  Always
-                installed *in addition to* the basics required by TinyAgent
+            log_manager: Log manager instance
+            default_python_codes: List of Python code snippets to execute before user code
+            code_tools: List of code tools to make available
+            pip_packages: List of pip packages to install in the sandbox
+            default_packages: List of default pip packages to install in the sandbox
+            apt_packages: List of apt packages to install in the sandbox
+            python_version: Python version to use in the sandbox
+            authorized_imports: Optional allow-list of modules the user code is permitted to import
+            authorized_functions: Optional allow-list of dangerous functions the user code is permitted to use
+            modal_secrets: Dictionary of secrets to make available to the sandbox
+            lazy_init: Whether to initialize Modal app lazily
+            sandbox_name: Name of the Modal sandbox
+            local_execution: Whether to execute code locally
+            check_string_obfuscation: If True (default), check for string obfuscation techniques. Set to False to allow legitimate use of base64 encoding and other string manipulations.
+            bypass_shell_safety: If True, bypass shell command safety checks (default: False for modal)
+            additional_safe_shell_commands: Additional shell commands to consider safe
+            additional_safe_control_operators: Additional shell control operators to consider safe
+            **kwargs: Additional keyword arguments
+        
+        Note:
+            The Modal sandbox is a secure environment for executing untrusted code.
+            It provides isolation from the host system and other sandboxes.
+            
+            Default packages are always installed, while pip_packages are added to
                 (git, curl, …) so you only need to specify the extras.
             python_version: Python version used for the sandbox image. If
                 ``None`` the current interpreter version is used.
@@ -63,7 +97,7 @@ class ModalProvider(CodeExecutionProvider):
             ]
 
         if apt_packages is None:
-            apt_packages = ["git", "curl", "nodejs", "npm"]
+            apt_packages = ["git", "curl", "nodejs", "npm","ripgrep"]
 
         if python_version is None:
             python_version = self.PYTHON_VERSION
@@ -74,6 +108,8 @@ class ModalProvider(CodeExecutionProvider):
         self.python_version: str = python_version
         self.authorized_imports = authorized_imports
 
+        self.authorized_functions = authorized_functions or []
+        self.check_string_obfuscation = check_string_obfuscation
         # ----------------------------------------------------------------------
         final_packages = list(set(self.default_packages + (pip_packages or [])))
         
@@ -84,6 +120,9 @@ class ModalProvider(CodeExecutionProvider):
             pip_packages=final_packages,
             secrets=modal_secrets or {},
             lazy_init=lazy_init,
+            bypass_shell_safety=bypass_shell_safety,
+            additional_safe_shell_commands=additional_safe_shell_commands,
+            additional_safe_control_operators=additional_safe_control_operators,
             **kwargs
         )
         
@@ -92,6 +131,7 @@ class ModalProvider(CodeExecutionProvider):
         self.modal_secrets = modal.Secret.from_dict(self.secrets)
         self.app = None
         self._app_run_python = None
+        self._app_run_shell = None
         self.is_trusted_code = kwargs.get("trust_code", False)
         
         self._setup_modal_app()
@@ -117,6 +157,7 @@ class ModalProvider(CodeExecutionProvider):
         )
         
         self._app_run_python = self.app.function()(_run_python)
+        self._app_run_shell = self.app.function()(_run_shell)
         
         # Add tools if provided
         if self.code_tools:
@@ -139,7 +180,7 @@ class ModalProvider(CodeExecutionProvider):
         full_code = "\n".join(code_lines)
         
         print("#" * 100)
-        print("#########################code#########################")
+        print("##########################################code##########################################")
         print(full_code)
         print("#" * 100)
 
@@ -170,6 +211,91 @@ class ModalProvider(CodeExecutionProvider):
         
         return clean_response(response)
     
+    async def execute_shell(
+        self,
+        command: List[str],
+        timeout: int = 30,
+        workdir: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Execute a shell command securely using Modal.
+        
+        Args:
+            command: List of command parts to execute
+            timeout: Maximum execution time in seconds
+            workdir: Working directory for command execution
+            
+        Returns:
+            Dictionary containing execution results with keys:
+            - stdout: stdout from the execution
+            - stderr: stderr from the execution
+            - exit_code: exit code from the command
+        """
+        # First, check if the command is safe to execute
+        timeout = min(timeout, self.TIMEOUT_MAX)
+        if type(command) == str:
+            command = command.split(" ")
+
+        print("#########################<Bash>#########################")
+        print(f"{COLOR['BLUE']}>{command}{COLOR['ENDC']}")
+        safety_check = self.is_safe_command(command)
+        if not safety_check["safe"]:
+            
+            response = {
+                "stdout": "",
+                "stderr": f"Command rejected for security reasons: {safety_check.get('reason', 'Unsafe command')}",
+                "exit_code": 1
+            }
+            print(f"{COLOR['RED']}{response['stderr']}{COLOR['ENDC']}")
+            return response
+        #execution_mode = "🏠 LOCALLY" if self.local_execution else "☁️ REMOTELY"
+        #print(f"Executing shell command {execution_mode} via Modal: {' '.join(command)}")
+        
+        # Show working directory information
+        if workdir:
+            print(f"Working directory: {workdir}")
+        
+        # If using Modal for remote execution
+        if not self.local_execution:
+            try:
+                with self.app.run():
+                    result = self._app_run_shell.remote(
+                        command=command,
+                        timeout=timeout,
+                        workdir=workdir
+                    )
+                
+                
+                print(f"{COLOR['GREEN']}{result}{COLOR['ENDC']}")
+                return result
+            except Exception as e:
+                response = {
+                    "stdout": "",
+                    "stderr": f"Error executing shell command: {str(e)}",
+                    "exit_code": 1
+                }
+                
+                print(f"{COLOR['RED']}{response['stderr']}{COLOR['ENDC']}")
+                return response
+        # If executing locally
+        else:
+            try:
+                result = self._app_run_shell.local(
+                    command=command,
+                    timeout=timeout,
+                    workdir=workdir
+                )
+                print(f"{COLOR['GREEN']}{result}{COLOR['ENDC']}")
+                return result
+            except Exception as e:
+                response = {
+                    "stdout": "",
+                    "stderr": f"Error executing shell command: {str(e)}",
+                    "exit_code": 1
+                }
+                print(f"{COLOR['RED']}{response['stderr']}{COLOR['ENDC']}")
+                return response
+    
     def _python_executor(self, code: str, globals_dict: Dict[str, Any] = None, locals_dict: Dict[str, Any] = None):
         """Execute Python code using Modal's native .local() or .remote() methods."""
         execution_mode = "🏠 LOCALLY" if self.local_execution else "☁️ REMOTELY"
@@ -191,8 +317,10 @@ class ModalProvider(CodeExecutionProvider):
                 full_code,
                 globals_dict or {},
                 locals_dict or {},
-                self.authorized_imports,
-                self.is_trusted_code,
+                authorized_imports=self.authorized_imports,
+                authorized_functions=self.authorized_functions,
+                trusted_code=self.is_trusted_code,
+                check_string_obfuscation=self.check_string_obfuscation,
             )
         else:
             with self.app.run():
@@ -200,8 +328,10 @@ class ModalProvider(CodeExecutionProvider):
                     full_code,
                     globals_dict or {},
                     locals_dict or {},
-                    self.authorized_imports,
-                    self.is_trusted_code,
+                    authorized_imports=self.authorized_imports,
+                    authorized_functions=self.authorized_functions,
+                    trusted_code=self.is_trusted_code,
+                    check_string_obfuscation=self.check_string_obfuscation,
                 )
     
     def _log_response(self, response: Dict[str, Any]):
@@ -224,14 +354,7 @@ class ModalProvider(CodeExecutionProvider):
             # Check if this is a security exception and highlight it in red if so
             error_text = response["error_traceback"]
             if "SECURITY" in error_text:
-                try:
-                    from ..modal_sandbox import COLOR
-                except ImportError:
-                    # Fallback colors if modal_sandbox is not available
-                    COLOR = {
-                        "RED": "\033[91m",
-                        "ENDC": "\033[0m",
-                    }
+                
                 print(f"{COLOR['RED']}{error_text}{COLOR['ENDC']}")
             else:
                 print(error_text)