From c66a4639b39951cf7af1d8b8e6d6cf05754ca0d2 Mon Sep 17 00:00:00 2001
From: Zhichao <Zhichao.lin@landing.ai>
Date: Thu, 10 Oct 2024 16:41:51 +0800
Subject: [PATCH] calling

---
 vision_agent/agent/vision_agent.py | 29 ++++++++++++++++-------------
 1 file changed, 16 insertions(+), 13 deletions(-)

diff --git a/vision_agent/agent/vision_agent.py b/vision_agent/agent/vision_agent.py
index ba6e1d64..ef49086a 100644
--- a/vision_agent/agent/vision_agent.py
+++ b/vision_agent/agent/vision_agent.py
@@ -21,7 +21,7 @@
     use_extra_vision_agent_args,
 )
 from vision_agent.utils import CodeInterpreterFactory
-from vision_agent.utils.execute import CodeInterpreter, Execution
+from vision_agent.utils.execute import CodeInterpreter, LocalCodeInterpreter, Execution
 
 logging.basicConfig(level=logging.INFO)
 _LOGGER = logging.getLogger(__name__)
@@ -29,6 +29,7 @@
 WORKSPACE.mkdir(parents=True, exist_ok=True)
 if str(WORKSPACE) != "":
     os.environ["PYTHONPATH"] = f"{WORKSPACE}:{os.getenv('PYTHONPATH', '')}"
+_SESSION_TIMEOUT = 600  # 10 minutes
 
 
 class BoilerplateCode:
@@ -195,7 +196,7 @@ def __init__(
         agent: Optional[LMM] = None,
         verbosity: int = 0,
         local_artifacts_path: Optional[Union[str, Path]] = None,
-        code_sandbox_runtime: Optional[str] = None,
+        code_interpreter: Optional[CodeInterpreter] = None,
         callback_message: Optional[Callable[[Dict[str, Any]], None]] = None,
     ) -> None:
         """Initialize the VisionAgent.
@@ -206,13 +207,17 @@ def __init__(
             verbosity (int): The verbosity level of the agent.
             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]): The code interpreter to use, default to local python
         """
 
         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
+            if code_interpreter is not None
+            else LocalCodeInterpreter(timeout=_SESSION_TIMEOUT)
+        )
         self.callback_message = callback_message
         if self.verbosity >= 1:
             _LOGGER.setLevel(logging.INFO)
@@ -284,9 +289,7 @@ 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:
+        with self.code_interpreter as code_interpreter:
             orig_chat = copy.deepcopy(chat)
             int_chat = copy.deepcopy(chat)
             last_user_message = chat[-1]
@@ -472,7 +475,7 @@ def __init__(
         agent: Optional[LMM] = None,
         verbosity: int = 0,
         local_artifacts_path: Optional[Union[str, Path]] = None,
-        code_sandbox_runtime: Optional[str] = None,
+        code_interpreter: Optional[CodeInterpreter] = None,
         callback_message: Optional[Callable[[Dict[str, Any]], None]] = None,
     ) -> None:
         """Initialize the VisionAgent using OpenAI LMMs.
@@ -483,7 +486,7 @@ def __init__(
             verbosity (int): The verbosity level of the agent.
             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]): The code interpreter to use, default to local python
         """
 
         agent = OpenAILMM(temperature=0.0, json_mode=True) if agent is None else agent
@@ -491,7 +494,7 @@ def __init__(
             agent,
             verbosity,
             local_artifacts_path,
-            code_sandbox_runtime,
+            code_interpreter,
             callback_message,
         )
 
@@ -502,7 +505,7 @@ def __init__(
         agent: Optional[LMM] = None,
         verbosity: int = 0,
         local_artifacts_path: Optional[Union[str, Path]] = None,
-        code_sandbox_runtime: Optional[str] = None,
+        code_interpreter: Optional[CodeInterpreter] = None,
         callback_message: Optional[Callable[[Dict[str, Any]], None]] = None,
     ) -> None:
         """Initialize the VisionAgent using Anthropic LMMs.
@@ -513,7 +516,7 @@ def __init__(
             verbosity (int): The verbosity level of the agent.
             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]): The code interpreter to use, default to local python
         """
 
         agent = AnthropicLMM(temperature=0.0) if agent is None else agent
@@ -521,6 +524,6 @@ def __init__(
             agent,
             verbosity,
             local_artifacts_path,
-            code_sandbox_runtime,
+            code_interpreter,
             callback_message,
         )