From 42455eeb6296f42bc2622f77ad0755e93f22bfe1 Mon Sep 17 00:00:00 2001
From: Asia <2736300+humpydonkey@users.noreply.github.com>
Date: Tue, 4 Jun 2024 23:19:09 -0700
Subject: [PATCH] Handle TimeoutError with retry (#113)

---
 vision_agent/utils/execute.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/vision_agent/utils/execute.py b/vision_agent/utils/execute.py
index 39df5bde..ad9468e8 100644
--- a/vision_agent/utils/execute.py
+++ b/vision_agent/utils/execute.py
@@ -426,6 +426,11 @@ def close(self, *args: Any, **kwargs: Any) -> None:
     def restart_kernel(self) -> None:
         self.interpreter.notebook.restart_kernel()
 
+    @tenacity.retry(
+        wait=tenacity.wait_exponential_jitter(),
+        stop=tenacity.stop_after_attempt(2),
+        retry=tenacity.retry_if_exception_type(TimeoutError),
+    )
     def exec_cell(self, code: str) -> Execution:
         execution = self.interpreter.notebook.exec_cell(code, timeout=self.timeout)
         return Execution.from_e2b_execution(execution)