diff --git a/vision_agent/agent/vision_agent_coder.py b/vision_agent/agent/vision_agent_coder.py index dd893d1d..b0b1c812 100644 --- a/vision_agent/agent/vision_agent_coder.py +++ b/vision_agent/agent/vision_agent_coder.py @@ -675,6 +675,7 @@ def chat_with_workflow( test_multi_plan: bool = True, display_visualization: bool = False, customized_tool_names: Optional[List[str]] = None, + envs: Optional[Dict[str, str]] = None, ) -> Dict[str, Any]: """Chat with VisionAgentCoder and return intermediate information regarding the task. @@ -692,6 +693,8 @@ def chat_with_workflow( show the image(s) created by visualization code (if there is any). customized_tool_names (List[str]): A list of customized tools for agent to pick and use. If not provided, default to full tool set from vision_agent.tools. + envs (Dict[str, str]): A dict of envs that would be passed to the CodeInterpreter. + Mainly for passing local envs to E2B. Local interpreter can pick up envs naturally. Returns: Dict[str, Any]: A dictionary containing the code, test, test result, plan, @@ -703,7 +706,7 @@ def chat_with_workflow( # NOTE: each chat should have a dedicated code interpreter instance to avoid concurrency issues with CodeInterpreterFactory.new_instance( - code_sandbox_runtime=self.code_sandbox_runtime + code_sandbox_runtime=self.code_sandbox_runtime, envs=envs ) as code_interpreter: chat = copy.deepcopy(chat) media_list = [] diff --git a/vision_agent/utils/execute.py b/vision_agent/utils/execute.py index 0de10335..d1c0b049 100644 --- a/vision_agent/utils/execute.py +++ b/vision_agent/utils/execute.py @@ -682,12 +682,13 @@ def get_default_instance() -> CodeInterpreter: def new_instance( code_sandbox_runtime: Optional[str] = None, remote_path: Optional[Union[str, Path]] = None, + envs: Optional[Dict[str, str]] = None, ) -> CodeInterpreter: if not code_sandbox_runtime: code_sandbox_runtime = os.getenv("CODE_SANDBOX_RUNTIME", "local") if code_sandbox_runtime == "e2b": instance: CodeInterpreter = E2BCodeInterpreter( - timeout=_SESSION_TIMEOUT, remote_path=remote_path + timeout=_SESSION_TIMEOUT, remote_path=remote_path, envs=envs ) elif code_sandbox_runtime == "local": instance = LocalCodeInterpreter(