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

tinyagent/code_agent/modal_sandbox.py

@@ -63,6 +63,7 @@ def create_sandbox(
     pip_install: Sequence[str] | None = None,
     image_name: str = "tinyagent-sandbox-image",
     app_name: str = "persistent-code-session",
+    force_build: bool = False,
     **sandbox_kwargs,
 ) -> Tuple[modal.Sandbox, modal.App]:
     """Create (or lookup) a `modal.Sandbox` pre-configured for code execution.
@@ -99,7 +100,7 @@ def create_sandbox(
 
     # Build image -----------------------------------------------------------
     agent_image = (
-        modal.Image.debian_slim(python_version=python_version)
+        modal.Image.debian_slim(python_version=python_version,force_build=force_build)
         .apt_install(*apt_packages)
         .pip_install(*full_pip_list)
     )
@@ -196,6 +197,7 @@ class SandboxSession:
         modal_secrets: modal.Secret,
         *,
         timeout: int = 5 * 60,
+        
         **create_kwargs,
     ) -> None:
         self.modal_secrets = modal_secrets

tinyagent/code_agent/providers/base.py

@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import Dict, List, Any, Optional
 from tinyagent.hooks.logging_manager import LoggingManager
+import cloudpickle
 
 
 class CodeExecutionProvider(ABC):
@@ -69,8 +70,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 +81,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 +104,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
@@ -149,4 +162,46 @@ 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 

tinyagent/code_agent/providers/modal_provider.py

@@ -26,6 +26,7 @@ class ModalProvider(CodeExecutionProvider):
         default_packages: Optional[List[str]] = None,
         apt_packages: Optional[List[str]] = None,
         python_version: Optional[str] = None,
+        authorized_imports: list[str] | None = None,
         modal_secrets: Dict[str, Union[str, None]] | None = None,
         lazy_init: bool = True,
         sandbox_name: str = "tinycodeagent-sandbox",
@@ -48,6 +49,7 @@ class ModalProvider(CodeExecutionProvider):
                 (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.
+            authorized_imports: Optional allow-list of modules the user code is permitted to import. Supports wildcard patterns (e.g. "pandas.*"). If ``None`` the safety layer blocks only the predefined dangerous modules.
         """
 
         # Resolve default values ------------------------------------------------
@@ -70,6 +72,7 @@ class ModalProvider(CodeExecutionProvider):
         self.default_packages: List[str] = default_packages
         self.apt_packages: List[str] = apt_packages
         self.python_version: str = python_version
+        self.authorized_imports = authorized_imports
 
         # ----------------------------------------------------------------------
         final_packages = list(set(self.default_packages + (pip_packages or [])))
@@ -89,6 +92,7 @@ class ModalProvider(CodeExecutionProvider):
         self.modal_secrets = modal.Secret.from_dict(self.secrets)
         self.app = None
         self._app_run_python = None
+        self.is_trusted_code = kwargs.get("trust_code", False)
         
         self._setup_modal_app()
         
@@ -139,17 +143,28 @@ class ModalProvider(CodeExecutionProvider):
         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)
         
@@ -164,18 +179,30 @@ class ModalProvider(CodeExecutionProvider):
         if self.executed_default_codes:
             print("✔️ default codes already executed")
             full_code = "\n".join(self.code_tools_definitions) +"\n\n"+code
+            # Code tools and default code are trusted, user code is not
         else:
             full_code = "\n".join(self.code_tools_definitions) +"\n\n"+ "\n".join(self.default_python_codes) + "\n\n" + code
             self.executed_default_codes = True
+            # First execution includes framework code which is trusted
         
         # Use Modal's native execution methods
         if self.local_execution:
-            # Use Modal's .local() method for local execution
-            return self._app_run_python.local(full_code, globals_dict or {}, locals_dict or {})
+            return self._app_run_python.local(
+                full_code,
+                globals_dict or {},
+                locals_dict or {},
+                self.authorized_imports,
+                self.is_trusted_code,
+            )
         else:
-            # Use Modal's .remote() method for remote execution
             with self.app.run():
-                return self._app_run_python.remote(full_code, globals_dict or {}, locals_dict or {})
+                return self._app_run_python.remote(
+                    full_code,
+                    globals_dict or {},
+                    locals_dict or {},
+                    self.authorized_imports,
+                    self.is_trusted_code,
+                )
     
     def _log_response(self, response: Dict[str, Any]):
         """Log the response from code execution."""
@@ -184,15 +211,31 @@ class ModalProvider(CodeExecutionProvider):
         print("#########################<printed_output>#########################")
         print(response["printed_output"])
         print("#########################</printed_output>#########################")
-        print("#########################<return_value>#########################")
-        print(response["return_value"])
-        print("#########################</return_value>#########################")
-        print("#########################<stderr>#########################")
-        print(response["stderr"])
-        print("#########################</stderr>#########################")
-        print("#########################<traceback>#########################")
-        print(response["error_traceback"])
-        print("#########################</traceback>#########################")
+        if response.get("return_value",None) not in [None,""]:
+            print("#########################<return_value>#########################")
+            print(response["return_value"])
+            print("#########################</return_value>#########################")
+        if response.get("stderr",None) not in [None,""]:
+            print("#########################<stderr>#########################")
+            print(response["stderr"])
+            print("#########################</stderr>#########################")
+        if response.get("error_traceback",None) not in [None,""]:
+            print("#########################<traceback>#########################")
+            # 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)
+            print("#########################</traceback>#########################")
     
     async def cleanup(self):
         """Clean up Modal resources."""