calling

vision_agent/agent/vision_agent.py

@@ -21,14 +21,15 @@
  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__)
 WORKSPACE = Path(os.getenv("WORKSPACE", ""))
 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,15 +486,15 @@ 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
  super().__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,14 +516,14 @@ 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
  super().__init__(
  agent,
  verbosity,
  local_artifacts_path,
- code_sandbox_runtime,
+ code_interpreter,
  callback_message,
  )