442 lines
13 KiB
Python
442 lines
13 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
多模式技能沙盒执行器
|
||
|
||
支持三种模式：
|
||
- none: 直接执行（默认，开发环境）
|
||
- docker: Docker 容器隔离
|
||
- kubernetes: Kubernetes Pod 隔离
|
||
|
||
环境变量：
|
||
SKILL_SANDBOX_MODE: 沙盒模式 (none/docker/kubernetes)，默认 none
|
||
SKILL_SANDBOX_IMAGE: Docker 镜像，默认 python:3.11-slim
|
||
SKILL_SANDBOX_MEMORY_LIMIT: 内存限制，默认 512m
|
||
SKILL_SANDBOX_CPU_LIMIT: CPU 限制，默认 1.0
|
||
SKILL_SANDBOX_NETWORK: 网络模式，默认 none
|
||
SKILL_SANDBOX_TIMEOUT: 超时时间（秒），默认 60
|
||
"""
|
||
|
||
import json
|
||
import logging
|
||
import os
|
||
from abc import ABC, abstractmethod
|
||
from typing import Any
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class SandboxBackend(ABC):
|
||
"""沙盒后端抽象基类"""
|
||
|
||
@abstractmethod
|
||
def execute(
|
||
self,
|
||
skill_name: str,
|
||
function_name: str,
|
||
function_args: dict,
|
||
) -> dict:
|
||
"""
|
||
执行技能函数
|
||
|
||
Args:
|
||
skill_name: 技能名称，如 "builtin/valuation_review"
|
||
function_name: 要执行的函数名，如 "build_dcf_report"
|
||
function_args: 函数参数字典
|
||
|
||
Returns:
|
||
执行结果字典
|
||
"""
|
||
pass
|
||
|
||
|
||
class NoSandboxBackend(SandboxBackend):
|
||
"""
|
||
无沙盒模式 - 直接执行（默认，仅用于开发环境）
|
||
|
||
特性：
|
||
- 直接导入并执行技能模块
|
||
- 零性能开销
|
||
- 无隔离，依赖代码审查保证安全
|
||
"""
|
||
|
||
# 函数名到脚本模块名的映射
|
||
FUNCTION_TO_SCRIPT_MAP = {
|
||
# valuation_review 技能
|
||
"build_dcf_report": "dcf_report",
|
||
"build_owner_earnings_report": "owner_earnings_report",
|
||
"build_ev_ebitda_report": "multiple_valuation_report",
|
||
"build_residual_income_report": "multiple_valuation_report",
|
||
}
|
||
|
||
def __init__(self):
|
||
self._module_cache = {}
|
||
|
||
def _get_script_name(self, function_name: str) -> str:
|
||
"""
|
||
根据函数名获取脚本模块名
|
||
|
||
优先使用预定义映射，否则尝试自动推断
|
||
"""
|
||
if function_name in self.FUNCTION_TO_SCRIPT_MAP:
|
||
return self.FUNCTION_TO_SCRIPT_MAP[function_name]
|
||
|
||
# 自动推断: build_X_report -> X_report
|
||
if function_name.startswith("build_") and function_name.endswith("_report"):
|
||
return function_name[6:] # 去掉 "build_" 前缀
|
||
|
||
return function_name
|
||
|
||
def execute(
|
||
self,
|
||
skill_name: str,
|
||
function_name: str,
|
||
function_args: dict,
|
||
) -> dict:
|
||
"""直接导入模块并执行函数"""
|
||
|
||
logger.debug(f"[NoSandbox] 执行技能: {skill_name}.{function_name}")
|
||
|
||
try:
|
||
# 将技能路径转换为模块路径
|
||
# builtin/valuation_review -> backend.skills.builtin.valuation_review.scripts
|
||
module_path = f"backend.skills.{skill_name.replace('/', '.')}.scripts"
|
||
|
||
# 从 function_name 获取脚本模块名
|
||
script_name = self._get_script_name(function_name)
|
||
submodule_path = f"{module_path}.{script_name}"
|
||
|
||
logger.debug(f"[NoSandbox] 导入模块: {submodule_path}.{function_name}")
|
||
|
||
# 缓存已加载的模块
|
||
if submodule_path not in self._module_cache:
|
||
self._module_cache[submodule_path] = __import__(
|
||
submodule_path,
|
||
fromlist=[function_name],
|
||
)
|
||
|
||
module = self._module_cache[submodule_path]
|
||
func = getattr(module, function_name)
|
||
|
||
# 执行函数
|
||
result = func(**function_args)
|
||
|
||
return {
|
||
"status": "success",
|
||
"result": result,
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(f"[NoSandbox] 执行失败: {e}")
|
||
return {
|
||
"status": "error",
|
||
"error": str(e),
|
||
"error_type": type(e).__name__,
|
||
}
|
||
|
||
|
||
class DockerSandboxBackend(SandboxBackend):
|
||
"""
|
||
Docker 沙盒模式 - 容器隔离
|
||
|
||
特性：
|
||
- 使用 Docker 容器隔离执行
|
||
- 支持资源限制（CPU、内存）
|
||
- 支持网络隔离
|
||
- 临时容器，执行后销毁
|
||
|
||
依赖：
|
||
pip install agentscope-runtime
|
||
Docker 守护进程运行中
|
||
"""
|
||
|
||
# 函数名到脚本模块名的映射
|
||
FUNCTION_TO_SCRIPT_MAP = {
|
||
# valuation_review 技能
|
||
"build_dcf_report": "dcf_report",
|
||
"build_owner_earnings_report": "owner_earnings_report",
|
||
"build_ev_ebitda_report": "multiple_valuation_report",
|
||
"build_residual_income_report": "multiple_valuation_report",
|
||
}
|
||
|
||
def __init__(self, config: dict):
|
||
self.config = config
|
||
self._available = None
|
||
|
||
def _get_script_name(self, function_name: str) -> str:
|
||
"""
|
||
根据函数名获取脚本模块名
|
||
|
||
优先使用预定义映射，否则尝试自动推断
|
||
"""
|
||
if function_name in self.FUNCTION_TO_SCRIPT_MAP:
|
||
return self.FUNCTION_TO_SCRIPT_MAP[function_name]
|
||
|
||
# 自动推断: build_X_report -> X_report
|
||
if function_name.startswith("build_") and function_name.endswith("_report"):
|
||
return function_name[6:] # 去掉 "build_" 前缀
|
||
|
||
return function_name
|
||
|
||
def _check_availability(self) -> bool:
|
||
"""检查 Docker 是否可用"""
|
||
if self._available is not None:
|
||
return self._available
|
||
|
||
try:
|
||
from agentscope_runtime.sandbox import BaseSandbox
|
||
self._available = True
|
||
except ImportError:
|
||
logger.error(
|
||
"AgentScope Runtime 未安装，无法使用 Docker 沙盒。"
|
||
"请运行: pip install agentscope-runtime"
|
||
)
|
||
self._available = False
|
||
|
||
return self._available
|
||
|
||
def execute(
|
||
self,
|
||
skill_name: str,
|
||
function_name: str,
|
||
function_args: dict,
|
||
) -> dict:
|
||
"""在 Docker 容器中执行"""
|
||
|
||
if not self._check_availability():
|
||
raise RuntimeError(
|
||
"Docker 沙盒不可用，请安装 agentscope-runtime "
|
||
"或切换到 SKILL_SANDBOX_MODE=none"
|
||
)
|
||
|
||
from agentscope_runtime.sandbox import BaseSandbox
|
||
|
||
logger.info(f"[DockerSandbox] 执行技能: {skill_name}.{function_name}")
|
||
|
||
# 获取脚本模块名
|
||
script_name = self._get_script_name(function_name)
|
||
|
||
# 构建执行代码
|
||
code = f"""
|
||
import sys
|
||
import json
|
||
|
||
# 挂载路径
|
||
sys.path.insert(0, '/skill/scripts')
|
||
|
||
# 导入函数
|
||
from {script_name} import {function_name}
|
||
|
||
# 执行
|
||
args = json.loads('{json.dumps(function_args)}')
|
||
result = {function_name}(**args)
|
||
|
||
# 输出结果
|
||
print(json.dumps({{"status": "success", "result": result}}))
|
||
"""
|
||
|
||
try:
|
||
with BaseSandbox(**self.config) as box:
|
||
# 挂载技能目录（只读）
|
||
host_skill_path = f"backend/skills/{skill_name}"
|
||
box.mount(
|
||
host_path=host_skill_path,
|
||
container_path="/skill",
|
||
read_only=True,
|
||
)
|
||
|
||
# 执行代码
|
||
exec_result = box.run_ipython_cell(code=code)
|
||
|
||
# 解析结果
|
||
if exec_result.get("exit_code") == 0:
|
||
output = exec_result.get("stdout", "")
|
||
return json.loads(output)
|
||
else:
|
||
return {
|
||
"status": "error",
|
||
"error": exec_result.get("stderr", "Unknown error"),
|
||
"exit_code": exec_result.get("exit_code"),
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(f"[DockerSandbox] 执行失败: {e}")
|
||
return {
|
||
"status": "error",
|
||
"error": str(e),
|
||
"error_type": type(e).__name__,
|
||
}
|
||
|
||
|
||
class KubernetesSandboxBackend(SandboxBackend):
|
||
"""
|
||
Kubernetes 沙盒模式 - Pod 隔离（预留接口）
|
||
|
||
特性：
|
||
- 使用 Kubernetes Pod 隔离执行
|
||
- 企业级隔离和调度
|
||
- 支持资源配额和命名空间
|
||
|
||
"""
|
||
|
||
def __init__(self, config: dict):
|
||
self.config = config
|
||
raise NotImplementedError(
|
||
"Kubernetes 沙盒模式尚未实现，"
|
||
"请使用 SKILL_SANDBOX_MODE=docker 或 none"
|
||
)
|
||
|
||
def execute(
|
||
self,
|
||
skill_name: str,
|
||
function_name: str,
|
||
function_args: dict,
|
||
) -> dict:
|
||
raise NotImplementedError()
|
||
|
||
|
||
class SkillSandbox:
|
||
"""
|
||
技能沙盒执行器
|
||
|
||
统一接口，根据配置自动选择后端。
|
||
默认使用 none 模式（无沙盒）。
|
||
|
||
示例：
|
||
>>> sandbox = SkillSandbox()
|
||
>>> result = sandbox.execute_skill(
|
||
... skill_name="builtin/valuation_review",
|
||
... function_name="build_dcf_report",
|
||
... function_args={"rows": [...], "current_date": "2024-01-01"}
|
||
... )
|
||
>>> print(result)
|
||
{"status": "success", "result": "..."}
|
||
"""
|
||
|
||
_instance = None
|
||
_mode = None
|
||
|
||
def __new__(cls):
|
||
"""单例模式"""
|
||
if cls._instance is None:
|
||
cls._instance = super().__new__(cls)
|
||
cls._instance._initialized = False
|
||
return cls._instance
|
||
|
||
def __init__(self):
|
||
if self._initialized:
|
||
return
|
||
|
||
self.mode = os.getenv("SKILL_SANDBOX_MODE", "none").lower()
|
||
self._backend = self._create_backend()
|
||
self._initialized = True
|
||
|
||
logger.debug(f"SkillSandbox 初始化完成，模式: {self.mode}")
|
||
|
||
def _create_backend(self) -> SandboxBackend:
|
||
"""根据模式创建对应后端"""
|
||
|
||
if self.mode == "none":
|
||
logger.debug("使用无沙盒模式（直接执行）")
|
||
return NoSandboxBackend()
|
||
|
||
elif self.mode == "docker":
|
||
config = {
|
||
"image": os.getenv(
|
||
"SKILL_SANDBOX_IMAGE", "python:3.11-slim"
|
||
),
|
||
"memory_limit": os.getenv(
|
||
"SKILL_SANDBOX_MEMORY_LIMIT", "512m"
|
||
),
|
||
"cpu_limit": float(
|
||
os.getenv("SKILL_SANDBOX_CPU_LIMIT", "1.0")
|
||
),
|
||
"network": os.getenv("SKILL_SANDBOX_NETWORK", "none"),
|
||
"timeout": int(os.getenv("SKILL_SANDBOX_TIMEOUT", "60")),
|
||
}
|
||
logger.info(f"使用 Docker 沙盒模式，配置: {config}")
|
||
return DockerSandboxBackend(config)
|
||
|
||
elif self.mode == "kubernetes":
|
||
config = {
|
||
"namespace": os.getenv(
|
||
"SKILL_SANDBOX_NAMESPACE", "agentscope"
|
||
),
|
||
"memory_limit": os.getenv(
|
||
"SKILL_SANDBOX_MEMORY_LIMIT", "512Mi"
|
||
),
|
||
"cpu_limit": os.getenv("SKILL_SANDBOX_CPU_LIMIT", "1000m"),
|
||
"timeout": int(os.getenv("SKILL_SANDBOX_TIMEOUT", "60")),
|
||
}
|
||
logger.info(f"使用 Kubernetes 沙盒模式，配置: {config}")
|
||
return KubernetesSandboxBackend(config)
|
||
|
||
else:
|
||
raise ValueError(
|
||
f"未知的沙盒模式: {self.mode}，"
|
||
f"请设置 SKILL_SANDBOX_MODE=none/docker/kubernetes"
|
||
)
|
||
|
||
def execute_skill(
|
||
self,
|
||
skill_name: str,
|
||
function_name: str,
|
||
function_args: dict | None = None,
|
||
) -> Any:
|
||
"""
|
||
执行技能函数
|
||
|
||
Args:
|
||
skill_name: 技能名称，如 "builtin/valuation_review"
|
||
function_name: 函数名，如 "build_dcf_report"
|
||
function_args: 函数参数，默认 None
|
||
|
||
Returns:
|
||
函数执行结果（成功时返回 result 字段，失败时抛出异常）
|
||
|
||
Raises:
|
||
RuntimeError: 执行失败
|
||
"""
|
||
if function_args is None:
|
||
function_args = {}
|
||
|
||
logger.debug(
|
||
f"执行技能: {skill_name}.{function_name} "
|
||
f"(模式: {self.mode})"
|
||
)
|
||
|
||
result = self._backend.execute(
|
||
skill_name=skill_name,
|
||
function_name=function_name,
|
||
function_args=function_args,
|
||
)
|
||
|
||
if result.get("status") == "error":
|
||
error_msg = result.get("error", "Unknown error")
|
||
error_type = result.get("error_type", "Exception")
|
||
raise RuntimeError(f"[{error_type}] {error_msg}")
|
||
|
||
return result.get("result")
|
||
|
||
@property
|
||
def current_mode(self) -> str:
|
||
"""获取当前沙盒模式"""
|
||
return self.mode
|
||
|
||
|
||
def get_sandbox() -> SkillSandbox:
|
||
"""
|
||
获取 SkillSandbox 单例实例
|
||
|
||
Returns:
|
||
SkillSandbox 实例
|
||
"""
|
||
return SkillSandbox()
|
||
|
||
|
||
def reset_sandbox():
|
||
"""
|
||
重置沙盒实例（用于测试）
|
||
"""
|
||
SkillSandbox._instance = None
|
||
SkillSandbox._mode = None
|