mirror of
https://github.com/langgenius/dify.git
synced 2026-01-14 06:07:33 +08:00
94 lines
3.1 KiB
Python
94 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from io import BytesIO
|
|
from types import TracebackType
|
|
|
|
from core.sandbox.bash_tool import SandboxBashTool
|
|
from core.sandbox.constants import DIFY_CLI_CONFIG_PATH, DIFY_CLI_PATH
|
|
from core.sandbox.dify_cli import DifyCliConfig
|
|
from core.sandbox.manager import SandboxManager
|
|
from core.session.cli_api import CliApiSessionManager
|
|
from core.tools.__base.tool import Tool
|
|
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SandboxSession:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
workflow_execution_id: str,
|
|
tenant_id: str,
|
|
user_id: str,
|
|
tools: list[Tool],
|
|
) -> None:
|
|
self._workflow_execution_id = workflow_execution_id
|
|
self._tenant_id = tenant_id
|
|
self._user_id = user_id
|
|
self._tools = tools
|
|
|
|
self._sandbox: VirtualEnvironment | None = None
|
|
self._bash_tool: SandboxBashTool | None = None
|
|
self._session_id: str | None = None
|
|
|
|
def __enter__(self) -> SandboxSession:
|
|
sandbox = SandboxManager.get(self._workflow_execution_id)
|
|
if sandbox is None:
|
|
raise RuntimeError(f"Sandbox not found for workflow_execution_id={self._workflow_execution_id}")
|
|
|
|
session = CliApiSessionManager().create(tenant_id=self._tenant_id, user_id=self._user_id)
|
|
self._session_id = session.id
|
|
|
|
try:
|
|
config = DifyCliConfig.create(session, self._tools)
|
|
config_json = json.dumps(config.model_dump(mode="json"), ensure_ascii=False)
|
|
|
|
sandbox.upload_file(DIFY_CLI_CONFIG_PATH, BytesIO(config_json.encode("utf-8")))
|
|
|
|
connection_handle = sandbox.establish_connection()
|
|
try:
|
|
future = sandbox.run_command(connection_handle, [DIFY_CLI_PATH, "init"])
|
|
result = future.result(timeout=30)
|
|
if result.is_error:
|
|
raise RuntimeError(f"Failed to initialize Dify CLI in sandbox: {result.error_message}")
|
|
finally:
|
|
sandbox.release_connection(connection_handle)
|
|
|
|
except Exception:
|
|
CliApiSessionManager().delete(session.id)
|
|
self._session_id = None
|
|
raise
|
|
|
|
self._sandbox = sandbox
|
|
self._bash_tool = SandboxBashTool(sandbox=sandbox, tenant_id=self._tenant_id)
|
|
return self
|
|
|
|
def __exit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc: BaseException | None,
|
|
tb: TracebackType | None,
|
|
) -> bool:
|
|
try:
|
|
self.cleanup()
|
|
except Exception:
|
|
logger.exception("Failed to cleanup SandboxSession")
|
|
return False
|
|
|
|
@property
|
|
def bash_tool(self) -> SandboxBashTool:
|
|
if self._bash_tool is None:
|
|
raise RuntimeError("SandboxSession is not initialized")
|
|
return self._bash_tool
|
|
|
|
def cleanup(self) -> None:
|
|
if self._session_id is None:
|
|
return
|
|
|
|
CliApiSessionManager().delete(self._session_id)
|
|
logger.debug("Cleaned up SandboxSession session_id=%s", self._session_id)
|
|
self._session_id = None
|