Enhance TinyCodeAgent and CodeExecutionProvider to synchronize user variables with execution context. This update ensures user variables are preserved and updated after code execution, even in the event of errors. Additionally, improved error handling for syntax errors during code parsing and enhanced global/local variable management in ModalProvider to maintain state consistency.

Author: InsightFactoryAPP

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 @@ def add_tools(self, tools: List[Any]) -> None:
         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,15 +81,29 @@ def add_tools(self, tools: List[Any]) -> None:
         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.
         
         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 @@ def get_user_variables(self) -> Dict[str, Any]:
         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

@@ -143,17 +143,28 @@ async def execute_python(self, code_lines: List[str], timeout: int = 120) -> Dic
         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)
 

tinyagent/code_agent/tiny_code_agent.py

@@ -261,7 +261,16 @@ def _setup_code_execution_tool(self):
         async def run_python(code_lines: List[str], timeout: int = 120) -> str:
             """Execute Python code using the configured provider."""
             try:
+                # Before execution, ensure provider has the latest user variables
+                if self.user_variables:
+                    self.code_provider.set_user_variables(self.user_variables)
+                    
                 result = await self.code_provider.execute_python(code_lines, timeout)
+                
+                # After execution, update TinyCodeAgent's user_variables from the provider
+                # This ensures they stay in sync
+                self.user_variables = self.code_provider.get_user_variables()
+                
                 return str(result)
             except Exception as e:
                 print("!"*100)
@@ -272,6 +281,11 @@ async def run_python(code_lines: List[str], timeout: int = 120) -> str:
                 print(f"{COLOR['RED']}{str(e)}{COLOR['ENDC']}")
                 print(f"{COLOR['RED']}{traceback.format_exc()}{COLOR['ENDC']}")
                 print("!"*100)
+                
+                # Even after an exception, update user_variables from the provider
+                # This ensures any variables that were successfully created/modified are preserved
+                self.user_variables = self.code_provider.get_user_variables()
+                
                 return f"Error executing code: {str(e)}"
 
         self.agent.add_tool(run_python)

tinyagent/code_agent/utils.py

@@ -116,22 +116,34 @@ def custom_print(*args, **kwargs):
     #updated_globals['print'] = custom_print
 
     # Parse the code
-    tree = ast.parse(code, mode="exec")
-    compiled = compile(tree, filename="<ast>", mode="exec")
+    try:
+        tree = ast.parse(code, mode="exec")
+        compiled = compile(tree, filename="<ast>", mode="exec")
+    except SyntaxError as e:
+        # Return syntax error without executing
+        return {
+            "printed_output": "", 
+            "return_value": None, 
+            "stderr": "", 
+            "error_traceback": f"Syntax error: {str(e)}",
+            "updated_globals": updated_globals,
+            "updated_locals": updated_locals
+        }
+    
     stdout_buf = io.StringIO()
     stderr_buf = io.StringIO()   
     # Execute with exception handling
     error_traceback = None
     output = None
 
+    # Merge all variables into globals to avoid scoping issues with generator expressions
+    # When exec() is called with both globals and locals, generator expressions can't
+    # access local variables. By using only globals, everything runs in global scope.
+    merged_globals = updated_globals.copy()
+    merged_globals.update(updated_locals)
+
     with contextlib.redirect_stdout(stdout_buf), contextlib.redirect_stderr(stderr_buf):
         try:
-            # Merge all variables into globals to avoid scoping issues with generator expressions
-            # When exec() is called with both globals and locals, generator expressions can't
-            # access local variables. By using only globals, everything runs in global scope.
-            merged_globals = updated_globals.copy()
-            merged_globals.update(updated_locals)
-            
             # Add 'exec' to authorized_functions for internal use
             internal_authorized_functions = ['exec','eval']
             if authorized_functions is not None and not isinstance(authorized_functions, bool):
@@ -152,6 +164,18 @@ def custom_print(*args, **kwargs):
         except Exception:
             # Capture the full traceback as a string
             error_traceback = traceback.format_exc()
+            
+            # CRITICAL FIX: Even when an exception occurs, we need to update the globals and locals
+            # with any variables that were successfully created/modified before the exception
+            for key, value in merged_globals.items():
+                # Skip special variables and modules
+                if key.startswith('__') or key in ['builtins', 'traceback', 'contextlib', 'io', 'ast', 'sys']:
+                    continue
+                    
+                # Update both dictionaries with the current state
+                if key in updated_locals or key not in updated_globals:
+                    updated_locals[key] = value
+                updated_globals[key] = value
 
     # Join all captured output
     #printed_output = ''.join(output_buffer)