diff --git a/taskweaver/ces/common.py b/taskweaver/ces/common.py index 81e2b85b..c59ac650 100644 --- a/taskweaver/ces/common.py +++ b/taskweaver/ces/common.py @@ -126,3 +126,7 @@ def get_session_client( cwd: Optional[str] = None, ) -> Client: ... + + @abstractmethod + def get_kernel_mode(self) -> Literal["local", "container"] | None: + ... diff --git a/taskweaver/ces/manager/sub_proc.py b/taskweaver/ces/manager/sub_proc.py index 88c84a68..da8678f2 100644 --- a/taskweaver/ces/manager/sub_proc.py +++ b/taskweaver/ces/manager/sub_proc.py @@ -64,12 +64,13 @@ def __init__( "TASKWEAVER_ENV_DIR", os.path.realpath(os.getcwd()), ) - if kernel_mode == "local": + self.kernel_mode = kernel_mode + if self.kernel_mode == "local": env_mode = EnvMode.Local - elif kernel_mode == "container": + elif self.kernel_mode == "container": env_mode = EnvMode.OutsideContainer else: - raise ValueError(f"Invalid kernel mode: {kernel_mode}, expected 'SubProcess' or 'Container'.") + raise ValueError(f"Invalid kernel mode: {self.kernel_mode}, expected 'local' or 'container'.") self.env = Environment( env_id, env_dir, @@ -98,3 +99,6 @@ def get_session_client( session_dir=session_dir, cwd=cwd, ) + + def get_kernel_mode(self) -> Literal["local", "container"] | None: + return self.kernel_mode diff --git a/taskweaver/chat/console/chat.py b/taskweaver/chat/console/chat.py index 205b3bd5..e22903be 100644 --- a/taskweaver/chat/console/chat.py +++ b/taskweaver/chat/console/chat.py @@ -406,6 +406,7 @@ class TaskWeaverChatApp(SessionEventHandlerBase): def __init__(self, app_dir: Optional[str] = None): self.app = TaskWeaverApp(app_dir=app_dir, use_local_uri=True) self.session = self.app.get_session() + self.exec_kernel_mode = self.session.code_executor.get_execution_mode() self.pending_files: List[Dict[Literal["name", "path", "content"], Any]] = [] def run(self): @@ -485,8 +486,18 @@ def _reset_session(self, first_session: bool = False): self._system_message("--- stopping the current session ---") self.session.stop() self.session = self.app.get_session() - self._system_message("--- new session starts ---") + self._system_message(f"--- new session starts in `{self.exec_kernel_mode}` mode ---") + if self.exec_kernel_mode == "local": + self._system_message( + "Code running in local mode " + "may incur security risks, such as file system access. " + "Please be cautious when executing code. " + "For higher security, consider using the `container` mode by setting " + "the `execution_service.kernel_mode` to `container`. " + "For more information, please refer to the documentation (" + "https://microsoft.github.io/TaskWeaver/docs/code_execution).", + ) self._assistant_message( "I am TaskWeaver, an AI assistant. To get started, could you please enter your request?", ) diff --git a/taskweaver/code_interpreter/code_executor.py b/taskweaver/code_interpreter/code_executor.py index 707442d9..a56524fa 100644 --- a/taskweaver/code_interpreter/code_executor.py +++ b/taskweaver/code_interpreter/code_executor.py @@ -57,6 +57,7 @@ def __init__( session_dir=workspace, cwd=execution_cwd, ) + self.exec_kernel_mode = self.exec_mgr.get_kernel_mode() self.client_started: bool = False self.plugin_registry = plugin_registry self.plugin_loaded: bool = False @@ -214,3 +215,6 @@ def format_code_output( lines.append("") return "\n".join([" " * indent + ln for ln in lines]) + + def get_execution_mode(self) -> Literal["local", "container"] | None: + return self.exec_kernel_mode