vision-agent 0.2.76__py3-none-any.whl → 0.2.77__py3-none-any.whl

vision_agent/utils/execute.py

@@ -11,10 +11,9 @@ import tempfile
 import traceback
 import warnings
 from enum import Enum
-from io import IOBase
 from pathlib import Path
 from time import sleep
-from typing import IO, Any, Dict, Iterable, List, Optional, Union, cast
+from typing import Any, Dict, Iterable, List, Optional, Union
 
 import nbformat
 import tenacity
@@ -33,6 +32,7 @@ from typing_extensions import Self
 
 load_dotenv()
 _LOGGER = logging.getLogger(__name__)
+_SESSION_TIMEOUT = 300  # 5 minutes
 
 
 class MimeType(str, Enum):
@@ -403,11 +403,8 @@ class CodeInterpreter(abc.ABC):
         self.restart_kernel()
         return self.exec_cell(code)
 
-    def upload_file(self, file: Union[str, Path, IO]) -> str:
+    def upload_file(self, file: Union[str, Path]) -> str:
         # Default behavior is a no-op (for local code interpreter)
-        assert not isinstance(
-            file, IO
-        ), "Don't pass IO objects to upload_file() of local interpreter"
         return str(file)
 
     def download_file(self, file_path: str) -> Path:
@@ -416,7 +413,6 @@ class CodeInterpreter(abc.ABC):
 
 
 class E2BCodeInterpreter(CodeInterpreter):
-    KEEP_ALIVE_SEC: int = 300
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
@@ -437,8 +433,8 @@ print(f"Vision Agent version: {va_version}")"""
         _LOGGER.info(f"E2BCodeInterpreter initialized:\n{sys_versions}")
 
     def close(self, *args: Any, **kwargs: Any) -> None:
-        self.interpreter.notebook.close()
         self.interpreter.close()
+        self.interpreter.kill()
 
     def restart_kernel(self) -> None:
         self.interpreter.notebook.restart_kernel()
@@ -449,25 +445,27 @@ print(f"Vision Agent version: {va_version}")"""
         retry=tenacity.retry_if_exception_type(TimeoutError),
     )
     def exec_cell(self, code: str) -> Execution:
-        self.interpreter.keep_alive(E2BCodeInterpreter.KEEP_ALIVE_SEC)
+        if not self.interpreter.is_running():
+            raise ConnectionResetError(
+                "Remote sandbox is closed unexpectedly. Please retry the operation."
+            )
+        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)
 
-    def upload_file(self, file: Union[str, Path, IO]) -> str:
-        try:
-            if isinstance(file, (Path, str)):
-                file = open(file, "rb")
-            return cast(str, self.interpreter.upload_file(cast(IO, file)))
-        finally:
-            assert isinstance(file, IOBase), f"Unexpected file type: {type(file)}"
-            file.close()
-            _LOGGER.info(f"File ({file}) is uploaded to: {file.name}")
+    def upload_file(self, file: Union[str, Path]) -> str:
+        file_name = Path(file).name
+        remote_path = f"/home/user/{file_name}"
+        with open(file, "rb") as f:
+            self.interpreter.files.write(path=remote_path, data=f)
+            _LOGGER.info(f"File ({file}) is uploaded to: {remote_path}")
+            return remote_path
 
     def download_file(self, file_path: str) -> Path:
-        file = tempfile.NamedTemporaryFile(mode="w+b", delete=False)
-        file.write(self.interpreter.download_file(file_path))
-        _LOGGER.info(f"File ({file_path}) is downloaded to: {file.name}")
-        return Path(file.name)
+        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)
 
     @staticmethod
     @tenacity.retry(
@@ -480,7 +478,7 @@ print(f"Vision Agent version: {va_version}")"""
 
 
 class LocalCodeInterpreter(CodeInterpreter):
-    def __init__(self, timeout: int = 600) -> None:
+    def __init__(self, timeout: int = _SESSION_TIMEOUT) -> None:
         super().__init__(timeout=timeout)
         self.nb = nbformat.v4.new_notebook()
         self.nb_client = NotebookClient(self.nb, timeout=self.timeout)
@@ -568,9 +566,9 @@ class CodeInterpreterFactory:
     @staticmethod
     def new_instance() -> CodeInterpreter:
         if os.getenv("CODE_SANDBOX_RUNTIME") == "e2b":
-            instance: CodeInterpreter = E2BCodeInterpreter(timeout=600)
+            instance: CodeInterpreter = E2BCodeInterpreter(timeout=_SESSION_TIMEOUT)
         else:
-            instance = LocalCodeInterpreter(timeout=600)
+            instance = LocalCodeInterpreter(timeout=_SESSION_TIMEOUT)
         atexit.register(instance.close)
         return instance
 

{vision_agent-0.2.76.dist-info → vision_agent-0.2.77.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vision-agent
-Version: 0.2.76
+Version: 0.2.77
 Summary: Toolset for Vision Agent
 Author: Landing AI
 Author-email: dev@landing.ai
@@ -10,7 +10,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Requires-Dist: e2b (>=0.17.1,<0.18.0)
-Requires-Dist: e2b-code-interpreter (>=0.0.9,<0.0.10)
+Requires-Dist: e2b-code-interpreter (==0.0.11a1)
 Requires-Dist: ipykernel (>=6.29.4,<7.0.0)
 Requires-Dist: langsmith (>=0.1.58,<0.2.0)
 Requires-Dist: moviepy (>=1.0.0,<2.0.0)

{vision_agent-0.2.76.dist-info → vision_agent-0.2.77.dist-info}/RECORD

@@ -12,12 +12,12 @@ vision_agent/tools/prompts.py,sha256=V1z4YJLXZuUl_iZ5rY0M5hHc_2tmMEUKr0WocXKGt4E
 vision_agent/tools/tool_utils.py,sha256=ZOY45bCX3nlo6iGwaZ8RVpRJB-vWxkXDed9oegT7-p0,1838
 vision_agent/tools/tools.py,sha256=TkZqNYX-ocwdaCdXd6c6tysSa_HX2y6Nrgl4JKni4IQ,43661
 vision_agent/utils/__init__.py,sha256=CW84HnhqI6XQVuxf2KifkLnSuO7EOhmuL09-gAymAak,219
-vision_agent/utils/execute.py,sha256=QImS69SN00logF-E68aNpT7YsJVRQOhZYlNLmCNEfro,21337
+vision_agent/utils/execute.py,sha256=DMaQz5-yULxDx-TlSMTRKOPHE7VmyR7PArhXXilm7h0,21368
 vision_agent/utils/image_utils.py,sha256=_cdiS5YrLzqkq_ZgFUO897m5M4_SCIThwUy4lOklfB8,7700
 vision_agent/utils/sim.py,sha256=ci6Eta73dDgLP1Ajtknbgmf1g8aAvBHqlVQvBuLMKXQ,4427
 vision_agent/utils/type_defs.py,sha256=BlI8ywWHAplC7kYWLvt4AOdnKpEW3qWEFm-GEOSkrFQ,1792
 vision_agent/utils/video.py,sha256=rNmU9KEIkZB5-EztZNlUiKYN0mm_55A_2VGUM0QpqLA,8779
-vision_agent-0.2.76.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-vision_agent-0.2.76.dist-info/METADATA,sha256=aJf41620hVI9F55vweMzP-pqb2x0iIHiyl5rZ-nUe2A,9438
-vision_agent-0.2.76.dist-info/WHEEL,sha256=7Z8_27uaHI_UZAc4Uox4PpBhQ9Y5_modZXWMxtUi4NU,88
-vision_agent-0.2.76.dist-info/RECORD,,
+vision_agent-0.2.77.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+vision_agent-0.2.77.dist-info/METADATA,sha256=2GjXWlij7wzd19pzbjNRt__AVhpVtLcAe_WfGnydxTI,9433
+vision_agent-0.2.77.dist-info/WHEEL,sha256=7Z8_27uaHI_UZAc4Uox4PpBhQ9Y5_modZXWMxtUi4NU,88
+vision_agent-0.2.77.dist-info/RECORD,,