Allow user to pass a template ID for E2B sandbox #170
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
WalkthroughThe update introduces an optional Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant CodeExecutionV2
participant Sandbox
Client->>CodeExecutionV2: Instantiate (with or without template)
CodeExecutionV2->>CodeExecutionV2: Initialize with template
CodeExecutionV2->>Sandbox: Create new instance with template
Sandbox-->>CodeExecutionV2: Return sandbox instance
CodeExecutionV2-->>Client: Return fully initialized instance
Poem
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (2)
athina/steps/code_execution_v2.py (2)
109-119: 🛠️ Refactor suggestionAdd template parameter to constructor.
The
templateattribute should be configurable during initialization.Apply this diff to add the template parameter:
def __init__( self, execution_environment: ExecutionEnvironment = EXECUTION_LOCAL, sandbox_timeout: Optional[int] = None, + template: Optional[str] = None, **data, ): super().__init__(**data) self._sandbox = None self.execution_environment = execution_environment self.sandbox_timeout = sandbox_timeout + self.template = template
231-324: 🛠️ Refactor suggestionReduce complexity of
_execute_e2bmethod.The method has a high cyclomatic complexity (13) and handles multiple responsibilities. Consider breaking it down into smaller, focused methods.
Apply this diff to refactor the method:
def _execute_e2b(self, input_data: dict, start_time: float) -> StepResult: try: self._create_or_initialize_sandbox() if self._sandbox is None: print("Sandbox is not initialized") return self._create_step_result( status="error", data="Sandbox is not initialized", start_time=start_time, ) - # Initialize input variables if we're running Python code - if not self.code.strip().startswith(COMMAND_PREFIX): - input_vars_code = self._prepare_input_variables(input_data) - if input_vars_code: - setup_code = "\n".join(input_vars_code) - setup_execution = self._sandbox.run_code(setup_code) - if setup_execution.error: - print( - f"Error setting up input variables: {setup_execution.error}" - ) - - # Execute code based on type (commands or Python) - if self.code.strip().startswith(COMMAND_PREFIX): - # Handle command execution - commands = [ - line.strip()[1:] for line in self.code.split("\n") if line.strip() - ] - output = [] - for command in commands: - command_result = self._sandbox.commands.run(command) - if command_result.error or command_result.exit_code != 0: - return self._create_step_result( - status="error", - data=f"Failed to execute command: {command}\nexit_code: {command_result.exit_code}\nDetails:\n{command_result.error}", - start_time=start_time, - ) - print(f"Command output: {command_result}") - if command_result.stdout: - output.extend(command_result.stdout) - return self._create_step_result( - status="success", - data="".join(output), - start_time=start_time, - exported_vars={}, - ) - else: - # Handle Python code execution - execution = self._sandbox.run_code(self.code) - if execution.error: - return self._create_step_result( - status="error", - data=f"Failed to execute the code.\nDetails:\n{execution.error}", - start_time=start_time, - ) - - # Capture variables for Python execution - var_execution = self._sandbox.run_code(VARIABLE_CAPTURE_CODE) - if var_execution.error: - print(f"Error capturing variables: {var_execution.error}") - return self._create_step_result( - status="success", - data="\n".join(execution.logs.stdout), - start_time=start_time, - exported_vars={}, - ) - - # Extract and return results - exported_vars = self._extract_exported_vars( - "\n".join(var_execution.logs.stdout) - ) - return self._create_step_result( - status="success", - data="\n".join(execution.logs.stdout), - start_time=start_time, - exported_vars=exported_vars, - ) + return ( + self._execute_commands(start_time) + if self.code.strip().startswith(COMMAND_PREFIX) + else self._execute_python_code(input_data, start_time) + ) except Exception as e: print(f"\nUnexpected error: {str(e)}") return self._create_step_result( status="error", data=f"Failed to execute the code.\nDetails:\n{str(e)}", start_time=start_time, ) + + def _execute_commands(self, start_time: float) -> StepResult: + """Execute shell commands in the sandbox.""" + commands = [ + line.strip()[1:] for line in self.code.split("\n") if line.strip() + ] + output = [] + for command in commands: + command_result = self._sandbox.commands.run(command) + if command_result.error or command_result.exit_code != 0: + return self._create_step_result( + status="error", + data=f"Failed to execute command: {command}\nexit_code: {command_result.exit_code}\nDetails:\n{command_result.error}", + start_time=start_time, + ) + print(f"Command output: {command_result}") + if command_result.stdout: + output.extend(command_result.stdout) + return self._create_step_result( + status="success", + data="".join(output), + start_time=start_time, + exported_vars={}, + ) + + def _execute_python_code(self, input_data: dict, start_time: float) -> StepResult: + """Execute Python code in the sandbox.""" + # Initialize input variables + if not self._initialize_input_variables(input_data): + return self._create_step_result( + status="error", + data="Failed to initialize input variables", + start_time=start_time, + ) + + # Execute code + execution = self._sandbox.run_code(self.code) + if execution.error: + return self._create_step_result( + status="error", + data=f"Failed to execute the code.\nDetails:\n{execution.error}", + start_time=start_time, + ) + + # Capture variables + exported_vars = self._capture_variables() + return self._create_step_result( + status="success", + data="\n".join(execution.logs.stdout), + start_time=start_time, + exported_vars=exported_vars, + ) + + def _initialize_input_variables(self, input_data: dict) -> bool: + """Initialize input variables in the sandbox.""" + input_vars_code = self._prepare_input_variables(input_data) + if input_vars_code: + setup_code = "\n".join(input_vars_code) + setup_execution = self._sandbox.run_code(setup_code) + if setup_execution.error: + print(f"Error setting up input variables: {setup_execution.error}") + return False + return True + + def _capture_variables(self) -> dict: + """Capture and extract variables from the sandbox.""" + var_execution = self._sandbox.run_code(VARIABLE_CAPTURE_CODE) + if var_execution.error: + print(f"Error capturing variables: {var_execution.error}") + return {} + return self._extract_exported_vars("\n".join(var_execution.logs.stdout))This refactoring:
- Splits the complex method into smaller, focused methods
- Improves error handling and readability
- Makes the code more maintainable and testable
🧰 Tools
🪛 GitHub Actions: Python Linter
[error] 231-231: 'CodeExecutionV2._execute_e2b' is too complex (13)
🧹 Nitpick comments (1)
athina/steps/code_execution_v2.py (1)
135-141: Add template validation.The template value should be validated before being passed to the Sandbox constructor to ensure it's not an empty string.
Apply this diff to add validation:
if self._sandbox is None: + if self.template is not None and not self.template.strip(): + raise ValueError("Template ID cannot be empty") self._sandbox = Sandbox( template=self.template, timeout=min( self.sandbox_timeout or self.DEFAULT_TIMEOUT, self.MAX_TIMEOUT ), metadata={"session_id": self.session_id}, )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
athina/steps/code_execution_v2.py(2 hunks)
🧰 Additional context used
🪛 GitHub Actions: Python Linter
athina/steps/code_execution_v2.py
[error] 231-231: 'CodeExecutionV2._execute_e2b' is too complex (13)
vivek-athina
approved these changes
Feb 13, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit