diff --git a/poetry.lock b/poetry.lock index eaa2d297..afb88f00 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. [[package]] name = "aenum" @@ -453,13 +453,13 @@ files = [ [[package]] name = "e2b" -version = "0.17.2a21" +version = "0.17.2a23" description = "E2B SDK that give agents cloud environments" optional = false python-versions = "<4.0,>=3.8" files = [ - {file = "e2b-0.17.2a21-py3-none-any.whl", hash = "sha256:2bb173a3c22e60690cc82ae163cc7fa8078824b57e9f7d1a5ab6970d8a86e235"}, - {file = "e2b-0.17.2a21.tar.gz", hash = "sha256:37d90745bc31096b00583fce0c4a17e801b2296d1a9f30a1399ad486f3587d83"}, + {file = "e2b-0.17.2a23-py3-none-any.whl", hash = "sha256:12e65f4f2c18c750c076f7bd22c24ed0b13ad0b59ecc9b99655990e7aeb7c93d"}, + {file = "e2b-0.17.2a23.tar.gz", hash = "sha256:948542564255d913da97a754765ea3ec1bf1482a897599ace9421e6857910cdd"}, ] [package.dependencies] @@ -474,17 +474,17 @@ urllib3 = ">=1.25.3" [[package]] name = "e2b-code-interpreter" -version = "0.0.11a1" +version = "0.0.11a2" description = "E2B Code Interpreter - Stateful code execution" optional = false python-versions = "<4.0,>=3.8" files = [ - {file = "e2b_code_interpreter-0.0.11a1-py3-none-any.whl", hash = "sha256:a212d8a1159f16fa6eef66c69c3570e6dc50046a1fa9435acdea919c3c9c7e32"}, - {file = "e2b_code_interpreter-0.0.11a1.tar.gz", hash = "sha256:5dad948a092bde56d91c83c9c8fb065845d40843e025077bfaa34f91afe3a5fd"}, + {file = "e2b_code_interpreter-0.0.11a2-py3-none-any.whl", hash = "sha256:bc3ae199acfb8e4c6675c4fc57a1070bb767a1e9bf99ac6d3aa09347fdc29997"}, + {file = "e2b_code_interpreter-0.0.11a2.tar.gz", hash = "sha256:6280f553be373865c43b2bfb546a3953d3faa43f9ce4ab7ed2dba576db548301"}, ] [package.dependencies] -e2b = "0.17.2a21" +e2b = "0.17.2a23" pydantic = "*" requests = ">=2.32.3,<3.0.0" websocket-client = ">=1.7.0,<2.0.0" @@ -3246,4 +3246,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools", [metadata] lock-version = "2.0" python-versions = ">=3.9,<4.0" -content-hash = "a2cf4a9da640d1984bcad58336a8e395851a8fd26020fb8e794ddef9541a94e0" +content-hash = "7dbe864df7d0b79afdcdc35b97209abb0f4a471e728f7929104861ededa2e12f" diff --git a/pyproject.toml b/pyproject.toml index 4f5fb2a7..e6451875 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,7 @@ rich = "^13.7.1" langsmith = "^0.1.58" ipykernel = "^6.29.4" e2b = "^0.17.1" -e2b-code-interpreter = "0.0.11a1" +e2b-code-interpreter = "0.0.11a2" tenacity = "^8.3.0" pillow-heif = "^0.16.0" pytube = "15.0.0" diff --git a/vision_agent/agent/vision_agent.py b/vision_agent/agent/vision_agent.py index 727178c7..d7886178 100644 --- a/vision_agent/agent/vision_agent.py +++ b/vision_agent/agent/vision_agent.py @@ -471,6 +471,7 @@ def __init__( tool_recommender: Optional[Sim] = None, verbosity: int = 0, report_progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None, + code_sandbox_runtime: Optional[str] = None, ) -> None: """Initialize the Vision Agent. @@ -487,6 +488,11 @@ def __init__( This is useful for streaming logs in a web application where multiple VisionAgent instances are running in parallel. This callback ensures that the progress are not mixed up. + code_sandbox_runtime: the code sandbox runtime to use. A code sandbox is + used to run the generated code. It can be one of the following + values: None, "local" or "e2b". If None, Vision Agent will read the + value from the environment variable CODE_SANDBOX_RUNTIME. If it's + also None, the local python runtime environment will be used. """ self.planner = ( @@ -506,6 +512,7 @@ def __init__( self.verbosity = verbosity self.max_retries = 2 self.report_progress_callback = report_progress_callback + self.code_sandbox_runtime = code_sandbox_runtime def __call__( self, @@ -560,7 +567,9 @@ def chat_with_workflow( raise ValueError("Chat cannot be empty.") # NOTE: each chat should have a dedicated code interpreter instance to avoid concurrency issues - with CodeInterpreterFactory.new_instance() as code_interpreter: + with CodeInterpreterFactory.new_instance( + code_sandbox_runtime=self.code_sandbox_runtime + ) as code_interpreter: chat = copy.deepcopy(chat) media_list = [] for chat_i in chat: diff --git a/vision_agent/utils/exceptions.py b/vision_agent/utils/exceptions.py new file mode 100644 index 00000000..1ffd3e11 --- /dev/null +++ b/vision_agent/utils/exceptions.py @@ -0,0 +1,42 @@ +"""Vision Agent exceptions.""" + + +class InvalidApiKeyError(Exception): + """Exception raised when the an invalid API key is provided. This error could be raised from any SDK code, not limited to a HTTP client.""" + + def __init__(self, message: str): + self.message = f"""{message} +For more information, see https://landing-ai.github.io/landingai-python/landingai.html#manage-api-credentials""" + super().__init__(self.message) + + def __str__(self) -> str: + return self.message + + +class RemoteSandboxError(Exception): + """Exception related to remote sandbox.""" + + is_retryable = False + + +class RemoteSandboxCreationError(RemoteSandboxError): + """Exception raised when failed to create a remote sandbox. + This could be due to the remote sandbox service is unavailable. + """ + + is_retryable = False + + +class RemoteSandboxExecutionError(RemoteSandboxError): + """Exception raised when failed in a remote sandbox code execution.""" + + is_retryable = False + + +class RemoteSandboxClosedError(RemoteSandboxError): + """Exception raised when a remote sandbox is dead. + This is retryable in the sense that the user can try again with a new sandbox. Can't be retried in the same sandbox. + When this error is raised, the user should retry by create a new VisionAgent (i.e. a new sandbox). + """ + + is_retryable = True diff --git a/vision_agent/utils/execute.py b/vision_agent/utils/execute.py index 46b6afd0..e7526384 100644 --- a/vision_agent/utils/execute.py +++ b/vision_agent/utils/execute.py @@ -1,5 +1,4 @@ import abc -import atexit import base64 import copy import logging @@ -18,7 +17,6 @@ import nbformat import tenacity from dotenv import load_dotenv -from e2b.api.v2.client.exceptions import ServiceException from e2b_code_interpreter import CodeInterpreter as E2BCodeInterpreterImpl from e2b_code_interpreter import Execution as E2BExecution from e2b_code_interpreter import Result as E2BResult @@ -30,9 +28,15 @@ from pydantic import BaseModel, field_serializer from typing_extensions import Self +from vision_agent.utils.exceptions import ( + RemoteSandboxClosedError, + RemoteSandboxCreationError, + RemoteSandboxExecutionError, +) + load_dotenv() _LOGGER = logging.getLogger(__name__) -_SESSION_TIMEOUT = 300 # 5 minutes +_SESSION_TIMEOUT = 600 # 10 minutes class MimeType(str, Enum): @@ -417,7 +421,15 @@ class E2BCodeInterpreter(CodeInterpreter): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) assert os.getenv("E2B_API_KEY"), "E2B_API_KEY environment variable must be set" - self.interpreter = E2BCodeInterpreter._new_e2b_interpreter_impl(*args, **kwargs) + try: + self.interpreter = E2BCodeInterpreter._new_e2b_interpreter_impl( + *args, **kwargs + ) + except Exception as e: + raise RemoteSandboxCreationError( + f"Failed to create a remote sandbox due to {e}" + ) from e + result = self.exec_cell( """ import platform @@ -433,27 +445,40 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: _LOGGER.info(f"E2BCodeInterpreter initialized:\n{sys_versions}") def close(self, *args: Any, **kwargs: Any) -> None: - self.interpreter.close() - self.interpreter.kill() + try: + self.interpreter.notebook.close() + self.interpreter.kill(request_timeout=2) + _LOGGER.info( + f"The sandbox {self.interpreter.sandbox_id} is closed successfully." + ) + except Exception as e: + _LOGGER.warn( + f"Failed to close the remote sandbox ({self.interpreter.sandbox_id}) due to {e}. This is not an issue. It's likely that the sandbox is already closed due to timeout." + ) def restart_kernel(self) -> None: + self._check_sandbox_liveness() self.interpreter.notebook.restart_kernel() @tenacity.retry( wait=tenacity.wait_exponential_jitter(), stop=tenacity.stop_after_attempt(2), + # TODO: change TimeoutError to a more specific exception when e2b team provides more granular retryable exceptions retry=tenacity.retry_if_exception_type(TimeoutError), ) def exec_cell(self, code: str) -> Execution: - if not self.interpreter.is_running(): - raise ConnectionResetError( - "Remote sandbox is closed unexpectedly. Please retry the operation." - ) + self._check_sandbox_liveness() self.interpreter.set_timeout(_SESSION_TIMEOUT) # Extend the life of the sandbox - execution = self.interpreter.notebook.exec_cell(code, timeout=self.timeout) - return Execution.from_e2b_execution(execution) + try: + execution = self.interpreter.notebook.exec_cell(code, timeout=self.timeout) + return Execution.from_e2b_execution(execution) + except Exception as e: + raise RemoteSandboxExecutionError( + f"Failed executing code in remote sandbox due to {e}: {code}" + ) 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: @@ -462,17 +487,26 @@ 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}") return Path(file.name) + def _check_sandbox_liveness(self) -> None: + try: + alive = self.interpreter.is_running(request_timeout=2) + except Exception as e: + _LOGGER.error( + f"Failed to check the health of the remote sandbox ({self.interpreter.sandbox_id}) due to {e}. Consider the sandbox as dead." + ) + alive = False + if not alive: + raise RemoteSandboxClosedError( + "Remote sandbox is closed unexpectedly. Please start a new VisionAgent instance." + ) + @staticmethod - @tenacity.retry( - wait=tenacity.wait_exponential_jitter(), - stop=tenacity.stop_after_delay(60), - retry=tenacity.retry_if_exception_type(ServiceException), - ) def _new_e2b_interpreter_impl(*args, **kwargs) -> E2BCodeInterpreterImpl: # type: ignore return E2BCodeInterpreterImpl(template="va-sandbox", *args, **kwargs) @@ -564,12 +598,17 @@ def get_default_instance() -> CodeInterpreter: return instance @staticmethod - def new_instance() -> CodeInterpreter: - if os.getenv("CODE_SANDBOX_RUNTIME") == "e2b": + def new_instance(code_sandbox_runtime: Optional[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) - else: + elif code_sandbox_runtime == "local": instance = LocalCodeInterpreter(timeout=_SESSION_TIMEOUT) - atexit.register(instance.close) + else: + raise ValueError( + f"Unsupported code sandbox runtime: {code_sandbox_runtime}. Supported runtimes: e2b, local" + ) return instance diff --git a/vision_agent/utils/type_defs.py b/vision_agent/utils/type_defs.py index 0b54c08d..d296a7ed 100644 --- a/vision_agent/utils/type_defs.py +++ b/vision_agent/utils/type_defs.py @@ -1,6 +1,8 @@ from pydantic import Field, field_validator from pydantic_settings import BaseSettings +from vision_agent.utils.exceptions import InvalidApiKeyError + class LandingaiAPIKey(BaseSettings): """The API key of a user in a particular organization in LandingLens. @@ -34,15 +36,3 @@ class Config: env_prefix = "landingai_" case_sensitive = False extra = "ignore" - - -class InvalidApiKeyError(Exception): - """Exception raised when the an invalid API key is provided. This error could be raised from any SDK code, not limited to a HTTP client.""" - - def __init__(self, message: str): - self.message = f"""{message} -For more information, see https://landing-ai.github.io/landingai-python/landingai.html#manage-api-credentials""" - super().__init__(self.message) - - def __str__(self) -> str: - return self.message