From 7e5e78d1f425b95c1c195762dfd6f69073ba3dae Mon Sep 17 00:00:00 2001
From: Zhichao <yzld2002@gmail.com>
Date: Fri, 11 Oct 2024 10:54:42 +0800
Subject: [PATCH] refactor: move e2b code interpreter to vas (#264)

* calling

* remove e2bcodeinterpreter and codefactory

* remove e2b

* remove e2b from pyproject

* remove tenacity

* add VERBOSITY env var

* Revert "remove tenacity"

This reverts commit 1e786ff876c63784dc4bb6d1d41d19b446336f77.

* Revert "remove e2b from pyproject"

This reverts commit b93d62fa046317e302522bc1ad97ecdac4a72948.

* Revert "remove e2b"

This reverts commit b0104ffd3fddfb2c705b5aded7227fa102e9976a.

* Revert "remove e2bcodeinterpreter and codefactory"

This reverts commit c67e2250d7da93ab4dda6409c2aae65d4999652a.

* Revert "calling"

This reverts commit c66a4639b39951cf7af1d8b8e6d6cf05754ca0d2.

* add an optional code_interpreter
---
 vision_agent/agent/vision_agent.py | 14 +++++++++++---
 vision_agent/tools/meta_tools.py   |  5 +++--
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/vision_agent/agent/vision_agent.py b/vision_agent/agent/vision_agent.py
index ba6e1d64..1a38468f 100644
--- a/vision_agent/agent/vision_agent.py
+++ b/vision_agent/agent/vision_agent.py
@@ -197,6 +197,7 @@ def __init__(
         local_artifacts_path: Optional[Union[str, Path]] = None,
         code_sandbox_runtime: Optional[str] = None,
         callback_message: Optional[Callable[[Dict[str, Any]], None]] = None,
+        code_interpreter: Optional[CodeInterpreter] = None,
     ) -> None:
         """Initialize the VisionAgent.
 
@@ -207,12 +208,14 @@ def __init__(
             local_artifacts_path (Optional[Union[str, Path]]): The path to the local
                 artifacts file.
             code_sandbox_runtime (Optional[str]): The code sandbox runtime to use.
+            code_interpreter (Optional[CodeInterpreter]): if not None, use this CodeInterpreter
         """
 
         self.agent = AnthropicLMM(temperature=0.0) if agent is None else agent
         self.max_iterations = 12
         self.verbosity = verbosity
         self.code_sandbox_runtime = code_sandbox_runtime
+        self.code_interpreter = code_interpreter
         self.callback_message = callback_message
         if self.verbosity >= 1:
             _LOGGER.setLevel(logging.INFO)
@@ -284,9 +287,14 @@ def chat_with_code(
             # this is setting remote artifacts path
             artifacts = Artifacts(WORKSPACE / "artifacts.pkl")
 
-        with CodeInterpreterFactory.new_instance(
-            code_sandbox_runtime=self.code_sandbox_runtime,
-        ) as code_interpreter:
+        code_interpreter = (
+            self.code_interpreter
+            if self.code_interpreter is not None
+            else CodeInterpreterFactory.new_instance(
+                code_sandbox_runtime=self.code_sandbox_runtime,
+            )
+        )
+        with code_interpreter:
             orig_chat = copy.deepcopy(chat)
             int_chat = copy.deepcopy(chat)
             last_user_message = chat[-1]
diff --git a/vision_agent/tools/meta_tools.py b/vision_agent/tools/meta_tools.py
index 024d4230..c9fc7be0 100644
--- a/vision_agent/tools/meta_tools.py
+++ b/vision_agent/tools/meta_tools.py
@@ -28,6 +28,7 @@
 CURRENT_LINE = 0
 DEFAULT_WINDOW_SIZE = 100
 ZMQ_PORT = os.environ.get("ZMQ_PORT", None)
+VERBOSITY = os.environ.get("VERBOSITY", 0)
 
 
 def report_progress_callback(port: int, inp: Dict[str, Any]) -> None:
@@ -375,7 +376,7 @@ def detect_dogs(image_path: str):
             )
         )
     else:
-        agent = va.agent.VisionAgentCoder()
+        agent = va.agent.VisionAgentCoder(verbosity=int(VERBOSITY))
 
     fixed_chat: List[Message] = [{"role": "user", "content": chat, "media": media}]
     response = agent.chat_with_workflow(
@@ -438,7 +439,7 @@ def detect_dogs(image_path: str):
             return dogs
     """
 
-    agent = va.agent.VisionAgentCoder()
+    agent = va.agent.VisionAgentCoder(verbosity=int(VERBOSITY))
     if name not in artifacts:
         print(f"[Artifact {name} does not exist]")
         return f"[Artifact {name} does not exist]"