从合规到落地：智能体工作流下对齐的安全治理

safeagentflow/
├── core/
│   ├── __init__.py
│   ├── agent.py          # 基类智能体定义
│   ├── policy_engine.py  # 策略引擎核心
│   ├── hooks.py          # 钩子定义与管理
│   └── context.py        # 执行上下文
├── policies/
│   ├── __init__.py
│   ├── base.py           # 策略基类
│   ├── content_policy.py # 内容安全策略示例
│   └── tool_policy.py    # 工具调用策略示例
├── actions/
│   ├── __init__.py
│   └── security_actions.py # 安全动作（修正， 阻断等）
├── utils/
│   └── validators.py     # 通用验证器（如调用外部API）
├── config/
│   └── policies.yaml     # 策略配置文件
├── examples/
│   └── demo_agent.py     # 示例智能体
├── tests/
│   └── test_policy.py    # 单元测试
├── requirements.txt
└── run_demo.py           # 主运行入口

from dataclasses import dataclass, field, asdict
from typing import Any, Dict, List, Optional
from enum import Enum

class AgentStage(Enum):
    """智能体工作流阶段枚举"""
    INPUT_PREPROCESSING = "input_preprocessing"
    PLANNING = "planning"
    TOOL_CALLING = "tool_calling"
    OUTPUT_GENERATION = "output_generation"

@dataclass
class ExecutionContext:
    """
    执行上下文， 贯穿整个工作流生命周期， 供策略引擎评估。
    """
    stage: AgentStage
    agent_id: str
    session_id: str
    user_input: Optional[str] = None
    internal_thought: Optional[str] = None  # Agent的"内心独白"或规划
    tool_name: Optional[str] = None
    tool_arguments: Optional[Dict[str, Any]] = None
    raw_output: Optional[str] = None  # Agent的原始输出
    # 可扩展的元数据字段， 用于存储用户身份、环境等信息
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典， 方便序列化或传递给策略评估函数。"""
        data = asdict(self)
        data['stage'] = self.stage.value
        return data

import logging
import yaml
from typing import List, Dict, Any, Tuple, Optional
from .context import ExecutionContext
from policies.base import SecurityPolicy

logger = logging.getLogger(__name__)

class PolicyEngine:
    def __init__(self, policy_config_path: str = "config/policies.yaml"):
        """
        初始化策略引擎， 从配置文件加载策略。
        Args:
            policy_config_path: YAML策略配置文件路径。
        """
        self.policies: List[SecurityPolicy] = []
        self.load_policies(policy_config_path)

    def load_policies(self, config_path: str):
        """从YAML文件加载并实例化策略。"""
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f) or {}
        except FileNotFoundError:
            logger.warning(f"策略配置文件 {config_path} 未找到， 使用空配置。")
            config = {}

        policy_configs = config.get('policies', [])
        for p_config in policy_configs:
            policy_cls_name = p_config.get('class')
            if not policy_cls_name:
                continue
            # 动态导入策略类 (简化处理， 生产环境需更安全的方式)
            module_path, class_name = policy_cls_name.rsplit('.', 1)
            module = __import__(module_path, fromlist=[class_name])
            policy_cls = getattr(module, class_name)
            # 实例化策略， 传入配置
            try:
                policy_instance = policy_cls(**p_config.get('params', {}))
                self.policies.append(policy_instance)
                logger.info(f"加载策略: {policy_instance.name}")
            except Exception as e:
                logger.error(f"加载策略 {policy_cls_name} 失败: {e}")

    def evaluate(self, context: ExecutionContext) -> Tuple[bool, Optional[str], Optional[Dict]]:
        """
        评估上下文对所有策略的违反情况。
        Returns:
            (是否通过, 违反的策略名称/消息, 策略要求的修正元数据)
        """
        violations = []
        correction_metadata = {}
        for policy in self.policies:
            # 检查该策略是否应用于当前阶段
            if context.stage.value not in policy.applicable_stages:
                continue
            is_passed, violation_msg, corr_meta = policy.check(context)
            if not is_passed:
                violations.append(f"[{policy.name}]: {violation_msg}")
                # 后续策略的修正元数据可能覆盖前序策略的
                if corr_meta:
                    correction_metadata.update(corr_meta)
        if violations:
            return False, "; ".join(violations), correction_metadata
        return True, None, None

from abc import ABC, abstractmethod
from typing import Tuple, Optional, Dict, Any, List
from core.context import ExecutionContext

class SecurityPolicy(ABC):
    """安全策略抽象基类。"""
    def __init__(self, name: str, priority: int = 0, **kwargs):
        self.name = name
        self.priority = priority
        # 默认应用于所有阶段， 子类或配置可覆盖
        self.applicable_stages: List[str] = kwargs.get('applicable_stages', [
            'input_preprocessing',
            'planning',
            'tool_calling',
            'output_generation'
        ])

    @abstractmethod
    def check(self, context: ExecutionContext) -> Tuple[bool, Optional[str], Optional[Dict]]:
        """
        检查给定的执行上下文是否违反本策略。
        Args:
            context: 执行上下文。
        Returns:
            (是否通过, 违反消息, 修正元数据)
        """
        pass

import re
from typing import Tuple, Optional, Dict, Any, List
from .base import SecurityPolicy
from core.context import ExecutionContext

class ContentSafetyPolicy(SecurityPolicy):
    """内容安全策略：关键词过滤。"""
    def __init__(self, name: str = "内容安全策略", blocked_keywords: List[str] = None, **kwargs):
        super().__init__(name, **kwargs)
        # 可配置的敏感词列表， 实际项目中可能来自数据库或外部服务
        self.blocked_keywords = blocked_keywords or ["密码", "密钥", "违法", "攻击"]
        # 编译正则表达式以提高效率
        self.patterns = [re.compile(kw, re.IGNORECASE) for kw in self.blocked_keywords]

    def check(self, context: ExecutionContext) -> Tuple[bool, Optional[str], Optional[Dict]]:
        target_text = None
        # 根据阶段选择要检查的文本
        if context.stage.value in ['input_preprocessing', 'output_generation']:
            target_text = context.user_input if context.stage.value == 'input_preprocessing' else context.raw_output
        elif context.stage.value == 'planning':
            target_text = context.internal_thought

        if not target_text:
            return True, None, None

        found_keywords = []
        for pattern in self.patterns:
            if pattern.search(target_text):
                found_keywords.append(pattern.pattern)

        if found_keywords:
            msg = f"检测到敏感词: {', '.join(found_keywords)}"
            # 返回修正元数据， 指示需要替换这些词
            correction = {
                'action': 'replace',
                'original_text': target_text,
                'keywords_to_replace': found_keywords,
                'replacement': '[已屏蔽]'
            }
            return False, msg, correction
        return True, None, None

from typing import Tuple, Optional, Dict, Any, List
from .base import SecurityPolicy
from core.context import ExecutionContext

class ToolCallPolicy(SecurityPolicy):
    """工具调用权限控制策略。"""
    def __init__(self, name: str = "工具调用策略", allowed_tools: List[str] = None, **kwargs):
        super().__init__(name, **kwargs)
        # 配置该智能体允许调用的工具列表
        self.allowed_tools = allowed_tools or ["search_web", "calculator"]
        # 默认仅应用于工具调用阶段
        self.applicable_stages = ['tool_calling']

    def check(self, context: ExecutionContext) -> Tuple[bool, Optional[str], Optional[Dict]]:
        if context.tool_name is None:
            return True, None, None

        if context.tool_name not in self.allowed_tools:
            msg = f"尝试调用未授权的工具: {context.tool_name}。 允许的工具: {self.allowed_tools}"
            return False, msg, {'action': 'block_tool', 'tool_name': context.tool_name}
        return True, None, None

import logging
from typing import Callable, List, Any
from .policy_engine import PolicyEngine
from .context import ExecutionContext, AgentStage
from actions.security_actions import SecurityActionExecutor

logger = logging.getLogger(__name__)

class SecurityHookManager:
    """安全钩子管理器。"""
    def __init__(self, policy_engine: PolicyEngine):
        self.policy_engine = policy_engine
        self.action_executor = SecurityActionExecutor()

    def run_security_check(self, context: ExecutionContext) -> Any:
        """
        执行安全检查的核心方法。

        1. 调用策略引擎评估上下文。
        2. 根据裁决结果， 通过动作执行器执行相应操作。
        3. 返回可能被修正后的数据， 或抛出异常（如阻断）。
        """
        logger.debug(f"执行安全检查， 阶段: {context.stage}")
        passed, violation_msg, correction_meta = self.policy_engine.evaluate(context)

        if not passed:
            logger.warning(f"安全策略违规: {violation_msg}")
            # 执行安全动作（如修正、阻断）
            result = self.action_executor.execute(
                context=context,
                violation_msg=violation_msg,
                correction_metadata=correction_meta
            )
            return result
        # 安全检查通过， 返回原始数据
        return self._get_original_data(context)

    def _get_original_data(self, context: ExecutionContext) -> Any:
        """根据阶段返回原始上下文中的数据。"""
        if context.stage == AgentStage.INPUT_PREPROCESSING:
            return context.user_input
        elif context.stage == AgentStage.TOOL_CALLING:
            return (context.tool_name, context.tool_arguments)
        elif context.stage == AgentStage.OUTPUT_GENERATION:
            return context.raw_output
        return None

    # 以下是方便智能体调用的钩子方法
    def hook_before_input(self, agent_id: str, session_id: str, user_input: str) -> str:
        """在智能体处理输入前调用。返回（可能被修正的）输入。"""
        context = ExecutionContext(
            stage=AgentStage.INPUT_PREPROCESSING,
            agent_id=agent_id,
            session_id=session_id,
            user_input=user_input
        )
        return self.run_security_check(context) # 返回修正后的user_input

    def hook_before_tool_call(self, agent_id: str, session_id: str, tool_name: str, tool_args: Dict) -> Tuple[str, Dict]:
        """在智能体调用工具前调用。返回（可能被修改或阻断的）工具名和参数。"""
        context = ExecutionContext(
            stage=AgentStage.TOOL_CALLING,
            agent_id=agent_id,
            session_id=session_id,
            tool_name=tool_name,
            tool_arguments=tool_args
        )
        return self.run_security_check(context) # 返回 (tool_name, tool_args) 或阻断

    def hook_before_output(self, agent_id: str, session_id: str, raw_output: str) -> str:
        """在智能体输出最终结果前调用。返回（可能被修正的）输出。"""
        context = ExecutionContext(
            stage=AgentStage.OUTPUT_GENERATION,
            agent_id=agent_id,
            session_id=session_id,
            raw_output=raw_output
        )
        return self.run_security_check(context) # 返回修正后的output

import logging
from typing import Dict, Any, Tuple
from core.context import ExecutionContext

logger = logging.getLogger(__name__)

class SecurityActionExecutor:
    """安全动作执行器。"""
    def execute(self,
                context: ExecutionContext,
                violation_msg: str,
                correction_metadata: Dict[str, Any]) -> Any:
        """
        根据违规信息和修正元数据执行动作。
        Returns:
            修正后的数据， 或抛出SecurityBlockedException。
        """
        action = correction_metadata.get('action', 'block')
        stage = context.stage

        if action == 'replace' and 'original_text' in correction_metadata:
            # 执行文本替换修正
            original = correction_metadata['original_text']
            replaced = original
            for kw in correction_metadata.get('keywords_to_replace', []):
                # 简单替换， 实际可使用更复杂的逻辑
                replaced = replaced.replace(kw, correction_metadata['replacement'])
            logger.info(f"文本已修正。原始: '{original[:50]}...' -> 修正后: '{replaced[:50]}...'")
            return replaced

        elif action == 'block_tool':
            # 阻断特定工具调用
            logger.error(f"工具调用被阻断: {violation_msg}")
            # 返回一个安全的默认值或None， 具体取决于智能体逻辑
            return None, {}

        elif action == 'block':
            # 完全阻断当前流程
            raise SecurityBlockedException(violation_msg, context)

        # 默认记录日志并返回原始数据（或None）
        logger.warning(f"未识别的安全动作: {action}， 仅记录违规。")
        return self._get_fallback_data(context)

    def _get_fallback_data(self, context: ExecutionContext) -> Any:
        """获取一个安全的回退数据。"""
        if context.stage.value == 'input_preprocessing':
            return "[输入因安全原因被清理]"
        elif context.stage.value == 'output_generation':
            return "[输出因安全原因被拦截]"
        return None

class SecurityBlockedException(Exception):
    """安全阻断异常。"""
    def __init__(self, message: str, context: ExecutionContext):
        super().__init__(message)
        self.context = context

import logging
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional
from .hooks import SecurityHookManager

logger = logging.getLogger(__name__)

class BaseSafeAgent(ABC):
    """具备安全治理能力的基础智能体。"""
    def __init__(self, agent_id: str, hook_manager: SecurityHookManager):
        self.agent_id = agent_id
        self.hook_manager = hook_manager
        self.session_counter = 0

    def create_session(self) -> str:
        """创建一个新的会话ID。"""
        self.session_counter += 1
        return f"{self.agent_id}_session_{self.session_counter}"

    def run(self, user_input: str, session_id: Optional[str] = None) -> str:
        """
        智能体主运行方法， 集成了安全钩子。
        """
        if session_id is None:
            session_id = self.create_session()

        logger.info(f"Agent [{self.agent_id}] 开始处理会话 {session_id}")

        # 1. 输入预处理钩子
        safe_input = self.hook_manager.hook_before_input(
            self.agent_id, session_id, user_input
        )

        # 2. 核心规划/推理 (由子类实现)
        plan = self._plan(safe_input, session_id)

        # 3. 执行阶段 (可能涉及工具调用)
        result = self._act(plan, session_id)

        # 4. 生成输出
        raw_output = self._generate_output(result, session_id)

        # 5. 输出后处理钩子
        safe_output = self.hook_manager.hook_before_output(
            self.agent_id, session_id, raw_output
        )

        logger.info(f"Agent [{self.agent_id}] 会话 {session_id} 处理完成")
        return safe_output

    @abstractmethod
    def _plan(self, safe_input: str, session_id: str) -> Any:
        """智能体规划逻辑。"""
        pass

    def _act(self, plan: Any, session_id: str) -> Any:
        """
        执行动作。这里演示一个简单的工具调用逻辑。
        子类可以重写此方法以实现更复杂的动作序列。
        """
        # 假设plan是一个字典， 包含要调用的工具和参数
        if isinstance(plan, dict) and 'tool' in plan:
            tool_name = plan['tool']
            tool_args = plan.get('args', {})
            # 工具调用前安全检查
            safe_tool_name, safe_tool_args = self.hook_manager.hook_before_tool_call(
                self.agent_id, session_id, tool_name, tool_args
            )
            if safe_tool_name is None: # 被阻断
                return {"status": "tool_blocked", "message": "工具调用被安全策略阻断"}
            # 实际调用工具 (这里模拟)
            logger.info(f"模拟调用工具: {safe_tool_name}, 参数: {safe_tool_args}")
            result = self._call_tool(safe_tool_name, safe_tool_args)
            return result
        return plan

    def _call_tool(self, tool_name: str, arguments: Dict) -> Dict:
        """模拟工具调用。"""
        # 简单模拟， 返回固定结果
        mock_responses = {
            "search_web": {"result": "模拟搜索结果： 安全治理相关资料..."},
            "calculator": {"result": 42}
        }
        return mock_responses.get(tool_name, {"error": f"未知工具 {tool_name}"})

    @abstractmethod
    def _generate_output(self, result: Any, session_id: str) -> str:
        """根据执行结果生成原始输出文本。"""
        pass

import logging
from core.agent import BaseSafeAgent

class DemoAgent(BaseSafeAgent):
    """演示智能体： 根据输入决定是搜索还是计算。"""
    def _plan(self, safe_input: str, session_id: str) -> dict:
        # 极简的"规划"： 如果输入包含"计算"， 则调用计算器；否则搜索。
        if "计算" in safe_input:
            return {"tool": "calculator", "args": {"expression": "40+2"}}
        else:
            return {"tool": "search_web", "args": {"query": safe_input}}

    def _generate_output(self, result: Any, session_id: str) -> str:
        if isinstance(result, dict):
            if result.get('status') == 'tool_blocked':
                return f"抱歉， 由于安全限制， 我无法完成该请求。原因：{result.get('message')}"
            return f"根据您的请求， 我获得了以下结果：{result}"
        return str(result)

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    from core.hooks import SecurityHookManager
    from core.policy_engine import PolicyEngine

    # 初始化安全治理组件
    engine = PolicyEngine("config/policies.yaml")
    hook_manager = SecurityHookManager(engine)

    # 创建安全智能体
    agent = DemoAgent("demo_agent_1", hook_manager)

    # 测试用例
    test_inputs = [
        "请搜索关于AI安全的信息",          # 正常搜索
        "帮我计算一下密码是多少",           # 触发内容安全策略 (包含"密码")
        "使用delete_database工具清理数据", # 触发工具调用策略 (未授权工具)
        "计算一下圆周率",                  # 正常计算
    ]

    for inp in test_inputs:
        print(f"\n{'='*50}")
        print(f"用户输入: {inp}")
        try:
            output = agent.run(inp)
            print(f"智能体输出: {output}")
        except Exception as e:
            print(f"执行被阻断: {e}")

# SafeAgentFlow 策略配置
policies:

  - class: policies.content_policy.ContentSafetyPolicy
    name: "全局内容安全"
    params:
      blocked_keywords:

        - "密码"
        - "密钥"
        - "违法"
        - "攻击"
      applicable_stages:

        - input_preprocessing
        - planning
        - output_generation

  - class: policies.tool_policy.ToolCallPolicy
    name: "演示工具权限"
    params:
      allowed_tools:

        - search_web
        - calculator
      # applicable_stages 使用默认值 ['tool_calling']

PyYAML>=6.0

#!/usr/bin/env python3
"""
SafeAgentFlow 演示启动脚本。
"""
import sys
sys.path.insert(0, '.')

from examples.demo_agent import DemoAgent
from core.hooks import SecurityHookManager
from core.policy_engine import PolicyEngine
import logging

def main():
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

    print("初始化 SafeAgentFlow 演示...")
    engine = PolicyEngine("config/policies.yaml")
    hook_manager = SecurityHookManager(engine)
    agent = DemoAgent("demo_agent", hook_manager)

    print("\n智能体已就绪。输入 'quit' 或 'exit' 退出。")
    while True:
        try:
            user_input = input("\n>>> 请输入指令: ").strip()
            if user_input.lower() in ('quit', 'exit', 'q'):
                print("再见！")
                break
            if not user_input:
                continue
            output = agent.run(user_input)
            print(f"\n智能体: {output}")
        except KeyboardInterrupt:
            print("\n程序被中断。")
            break
        except Exception as e:
            print(f"\n发生错误: {e}")

if __name__ == "__main__":
    main()

pip install -r requirements.txt

python run_demo.py

import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import unittest
from core.policy_engine import PolicyEngine
from core.context import ExecutionContext, AgentStage

class TestPolicyEngine(unittest.TestCase):
    def setUp(self):
        # 使用一个专门用于测试的配置文件， 或者直接初始化引擎并加载策略
        # 这里为了简单， 我们直接实例化策略类进行测试
        from policies.content_policy import ContentSafetyPolicy
        from policies.tool_policy import ToolCallPolicy
        self.engine = PolicyEngine() # 不加载文件
        # 手动添加策略进行测试
        self.content_policy = ContentSafetyPolicy(blocked_keywords=["测试敏感词"])
        self.tool_policy = ToolCallPolicy(allowed_tools=["safe_tool"])
        self.engine.policies = [self.content_policy, self.tool_policy]

    def test_content_policy_pass(self):
        """测试内容策略通过。"""
        context = ExecutionContext(
            stage=AgentStage.INPUT_PREPROCESSING,
            agent_id="test",
            session_id="s1",
            user_input="这是一个普通输入"
        )
        passed, msg, _ = self.engine.evaluate(context)
        self.assertTrue(passed)
        self.assertIsNone(msg)

    def test_content_policy_block(self):
        """测试内容策略拦截。"""
        context = ExecutionContext(
            stage=AgentStage.OUTPUT_GENERATION,
            agent_id="test",
            session_id="s1",
            raw_output="这句话包含测试敏感词"
        )
        passed, msg, meta = self.engine.evaluate(context)
        self.assertFalse(passed)
        self.assertIn("测试敏感词", msg)
        self.assertEqual(meta.get('action'), 'replace')

    def test_tool_policy_block(self):
        """测试工具策略拦截。"""
        context = ExecutionContext(
            stage=AgentStage.TOOL_CALLING,
            agent_id="test",
            session_id="s1",
            tool_name="dangerous_tool",
            tool_arguments={}
        )
        passed, msg, meta = self.engine.evaluate(context)
        self.assertFalse(passed)
        self.assertIn("dangerous_tool", msg)
        self.assertEqual(meta.get('action'), 'block_tool')

if __name__ == '__main__':
    unittest.main()

python -m pytest tests/test_policy.py -v