Add back code

vision_agent/utils/execute.py

@@ -462,6 +462,7 @@ def close(self, *args: Any, **kwargs: Any) -> None:
  )
 
  def restart_kernel(self) -> None:
+ self._check_sandbox_liveness()
  self.interpreter.notebook.restart_kernel()
 
  @tenacity.retry(
@@ -473,6 +474,7 @@ def restart_kernel(self) -> None:
  after=tenacity.after_log(_LOGGER, logging.INFO),
  )
  def exec_cell(self, code: str) -> Execution:
+ self._check_sandbox_liveness()
  self.interpreter.set_timeout(_SESSION_TIMEOUT) # Extend the life of the sandbox
  try:
  _LOGGER.info(f"Start code execution in remote sandbox {self.interpreter.sandbox_id}. Timeout: {_SESSION_TIMEOUT}. Code hash: {hash(code)}")
@@ -485,6 +487,7 @@ def exec_cell(self, code: str) -> Execution:
  ) from e
 
  def upload_file(self, file: Union[str, Path]) -> str:
+ self._check_sandbox_liveness()
  file_name = Path(file).name
  remote_path = f"/home/user/{file_name}"
  with open(file, "rb") as f:
@@ -493,6 +496,7 @@ def upload_file(self, file: Union[str, Path]) -> str:
  return remote_path
 
  def download_file(self, file_path: str) -> Path:
+ self._check_sandbox_liveness()
  with tempfile.NamedTemporaryFile(mode="w+b", delete=False) as file:
  file.write(self.interpreter.files.read(path=file_path, format="bytes"))
  _LOGGER.info(f"File ({file_path}) is downloaded to: {file.name}")
@@ -514,7 +518,7 @@ def _check_sandbox_liveness(self) -> None:
  @staticmethod
  def _new_e2b_interpreter_impl(*args, **kwargs) -> E2BCodeInterpreterImpl: # type: ignore
  template_name = os.environ.get("E2B_TEMPLATE_NAME", "va-sandbox")
- _LOGGER.info(f"Creating a new E2BCodeInterpreter instance using template: {template_name}")
+ _LOGGER.info(f"Creating a new E2BCodeInterpreter using template: {template_name}")
  return E2BCodeInterpreterImpl(template=template_name, *args, **kwargs)