tinyagent-py 0.0.12__py3-none-any.whl → 0.0.15__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/base.py

@@ -1,6 +1,7 @@
 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
 
 
 class CodeExecutionProvider(ABC):
@@ -34,6 +35,15 @@ class CodeExecutionProvider(ABC):
         self._locals_dict = kwargs.get("locals_dict", {})
         self._user_variables = {}
         self.code_tools_definitions = []
+        # 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"
+        }
+        # Safe control operators for shell commands
+        self.safe_control_operators: Set[str] = {"&&", "||", ";", "|"}
     
     @abstractmethod
     async def execute_python(
@@ -57,6 +67,129 @@ 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 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."""
@@ -69,8 +202,6 @@ class CodeExecutionProvider(ABC):
         Args:
             tools: List of tool objects to add
         """
-        import cloudpickle
-        
         tools_str_list = ["import cloudpickle"]
         tools_str_list.append("###########<tools>###########\n")
         for tool in tools:
@@ -82,6 +213,22 @@ class CodeExecutionProvider(ABC):
         tools_str_list.append("\n\n")
         self.code_tools_definitions.extend(tools_str_list)
     
+    def set_code_tools(self, tools: List[Any]) -> None:
+        """
+        Set the code tools available in the execution environment.
+        Replaces any existing tools with the new list.
+        
+        Args:
+            tools: List of tool objects to set
+        """
+        # Clear existing tools
+        self.code_tools = tools.copy()
+        self.code_tools_definitions = []
+        
+        # Add the new tools
+        if tools:
+            self.add_tools(tools)
+    
     def set_user_variables(self, variables: Dict[str, Any]) -> None:
         """
         Set user variables that will be available in the Python environment.
@@ -89,8 +236,6 @@ class CodeExecutionProvider(ABC):
         Args:
             variables: Dictionary of variable name -> value pairs
         """
-        import cloudpickle
-        
         self._user_variables = variables.copy()
         
         # Add variables to the execution environment by serializing them
@@ -116,14 +261,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."""
@@ -131,16 +276,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]:
         """
@@ -149,4 +294,60 @@ class CodeExecutionProvider(ABC):
         Returns:
             Dictionary of current user variables
         """
-        return self._user_variables.copy() 
+        return self._user_variables.copy()
+    
+    def update_user_variables_from_globals(self, globals_dict: Dict[str, Any]) -> None:
+        """
+        Extract and update user variables from the globals dictionary after code execution.
+        This ensures that any modifications to user variables during code execution are preserved.
+        
+        Args:
+            globals_dict: The globals dictionary after code execution
+        """
+        if not globals_dict or not self._user_variables:
+            return
+            
+        # Update user variables with values from globals
+        for var_name in list(self._user_variables.keys()):
+            if var_name in globals_dict:
+                try:
+                    # Try to serialize the value to ensure it's valid
+                    cloudpickle.dumps(globals_dict[var_name])
+                    # Update the user variable with the new value
+                    self._user_variables[var_name] = globals_dict[var_name]
+                except Exception:
+                    # If serialization fails, keep the old value
+                    pass
+                    
+        # Check for new variables that might have been created
+        # This handles cases where LLM creates new variables that should be preserved
+        for var_name, var_value in globals_dict.items():
+            # Skip special variables, modules, and functions
+            if (var_name.startswith('__') or 
+                var_name in ['builtins', 'cloudpickle'] or
+                callable(var_value) or
+                var_name in self._user_variables):
+                continue
+                
+            try:
+                # Try to serialize the value to ensure it's valid
+                cloudpickle.dumps(var_value)
+                # Add the new variable to user variables
+                self._user_variables[var_name] = var_value
+            except Exception:
+                # If serialization fails, skip this variable
+                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,39 @@ 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,
         **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.
+            **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 +91,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 +102,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 [])))
         
@@ -92,6 +122,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 +148,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,26 +171,122 @@ class ModalProvider(CodeExecutionProvider):
         full_code = "\n".join(code_lines)
         
         print("#" * 100)
-        print("#########################code#########################")
+        print("##########################################code##########################################")
         print(full_code)
         print("#" * 100)
 
-
         
         # Use Modal's native execution methods
         response = self._python_executor(full_code, self._globals_dict, self._locals_dict)
         
         print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!<response>!!!!!!!!!!!!!!!!!!!!!!!!!")
         
-        # Update the instance globals and locals with the execution results
-        self._globals_dict = cloudpickle.loads(make_session_blob(response["updated_globals"]))
-        self._locals_dict = cloudpickle.loads(make_session_blob(response["updated_locals"]))
-
+        # Always update globals and locals dictionaries, regardless of whether there was an error
+        # This ensures variables are preserved even when code execution fails
+        try:
+            # Update globals and locals from the response
+            if "updated_globals" in response:
+                self._globals_dict = cloudpickle.loads(make_session_blob(response["updated_globals"]))
+                
+            if "updated_locals" in response:
+                self._locals_dict = cloudpickle.loads(make_session_blob(response["updated_locals"]))
+                
+            # Update user variables from the updated globals and locals
+            # This preserves any changes made to variables by the LLM
+            self.update_user_variables_from_globals(self._globals_dict)
+            self.update_user_variables_from_globals(self._locals_dict)
+        except Exception as e:
+            print(f"Warning: Failed to update globals/locals after execution: {str(e)}")
         
         self._log_response(response)
         
         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"
@@ -180,8 +308,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():
@@ -189,8 +319,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]):
@@ -213,14 +345,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)

tinyagent/code_agent/safety.py

@@ -295,7 +295,8 @@ def _detect_string_obfuscation(tree: ast.AST) -> bool:
 
 
 def validate_code_safety(code: str, *, authorized_imports: Sequence[str] | None = None, 
-                        authorized_functions: Sequence[str] | None = None, trusted_code: bool = False) -> None:
+                        authorized_functions: Sequence[str] | None = None, trusted_code: bool = False,
+                        check_string_obfuscation: bool = True) -> None:
     """Static validation of user code.
 
     Parameters
@@ -312,6 +313,9 @@ def validate_code_safety(code: str, *, authorized_imports: Sequence[str] | None
     trusted_code
         If True, skip security checks. This should only be used for code that is part of the
         framework, developer-provided tools, or default executed code.
+    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.
     """
     # Skip security checks for trusted code
     if trusted_code:
@@ -384,7 +388,7 @@ def validate_code_safety(code: str, *, authorized_imports: Sequence[str] | None
     # ------------------------------------------------------------------
     # Detect string obfuscation techniques that might be used to bypass security
     # ------------------------------------------------------------------
-    if _detect_string_obfuscation(tree):
+    if check_string_obfuscation and _detect_string_obfuscation(tree):
         raise ValueError("SECURITY VIOLATION: Suspicious string manipulation detected that could be used to bypass security.")
 
     if blocked: