Agent 策略引擎设计：用代码定义 Agent 行为边界

from functools import wraps
from typing import Dict, Any, Callable
import datetime

class PolicyEngine:
    def __init__(self, policies: list):
        self.policies = policies
    
    def evaluate(self, principal: str, action: str, 
                 resource: Dict, context: Dict) -> bool:
        """评估策略，返回是否允许。forbid 优先于 permit。"""
        # 先检查所有 forbid 规则
        for policy in self.policies:
            if policy['effect'] == 'forbid' and \
               self._match(policy, principal, action, resource, context):
                return False
        # 再检查 permit 规则
        for policy in self.policies:
            if policy['effect'] == 'permit' and \
               self._match(policy, principal, action, resource, context):
                return True
        # 默认拒绝
        return False
    
    def _match(self, policy: Dict, principal: str, action: str,
               resource: Dict, context: Dict) -> bool:
        """匹配策略条件"""
        if policy.get('principal') and policy['principal'] != principal:
            return False
        if policy.get('action') and policy['action'] != action:
            return False
        if policy.get('resource_type') and \
           policy['resource_type'] != resource.get('type'):
            return False
        # 时间条件检查
        if 'time_before' in policy:
            current_hour = datetime.datetime.now().hour
            if current_hour >= policy['time_before']:
                return False
        # 语义条件（简化示例）
        if 'semantic_check' in policy:
            if not policy['semantic_check'](resource, context):
                return False
        return True


def policy_check(engine: PolicyEngine, principal: str):
    """装饰器：在工具调用前执行策略检查"""
    def decorator(func: Callable):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # 提取资源信息（简化示例）
            resource = {
                'type': kwargs.get('resource_type', 'unknown'),
                'id': kwargs.get('resource_id', 'unknown'),
                'owner': kwargs.get('owner', 'unknown')
            }
            context = {
                'session_id': kwargs.get('session_id'),
                'timestamp': datetime.datetime.now()
            }
            
            # 策略评估
            if not engine.evaluate(principal, func.__name__, resource, context):
                raise PermissionError(
                    f"策略拒绝: {principal} 尝试执行 {func.__name__} "
                    f"在 {resource['type']}:{resource['id']}"
                )
            return func(*args, **kwargs)
        return wrapper
    return decorator


# 定义策略
policies = [
    {
        'effect': 'forbid',
        'resource_type': 'database',
        'time_before': 9,  # 9点前禁止数据库操作
    },
    {
        'effect': 'permit',
        'principal': 'production-agent',
        'action': 'read_document',
        'resource_type': 'document',
    },
    {
        'effect': 'forbid',
        'principal': 'production-agent',
        'action': 'delete_*',  # 禁止所有删除操作
    },
]

engine = PolicyEngine(policies)

# 在工具上应用策略检查
@policy_check(engine, 'production-agent')
def read_document(resource_type: str, resource_id: str, 
                  owner: str = None, session_id: str = None):
    """读取文档"""
    # 实际工具逻辑
    return f"读取文档 {resource_id} 内容"


@policy_check(engine, 'production-agent')
def delete_document(resource_type: str, resource_id: str,
                    owner: str = None, session_id: str = None):
    """删除文档"""
    # 实际工具逻辑
    return f"删除文档 {resource_id}"


# 在 LangGraph 中集成
from langgraph.graph import StateGraph, END

class AgentState(dict):
    messages: list
    current_tool: str
    resource_info: dict

def tool_node(state: AgentState):
    """工具执行节点，策略检查在装饰器中自动完成"""
    tool_map = {
        'read_document': read_document,
        'delete_document': delete_document,
    }
    tool = tool_map[state['current_tool']]
    try:
        result = tool(**state['resource_info'])
        return {'messages': state['messages'] + [result]}
    except PermissionError as e:
        return {'messages': state['messages'] + [f"权限错误: {e}"]}

# 构建图
graph = StateGraph(AgentState)
graph.add_node("tool", tool_node)
graph.add_edge("tool", END)
app = graph.compile()

# 测试
try:
    read_document(
        resource_type='document', 
        resource_id='doc-123',
        owner='team-a',
        session_id='sess-1'
    )
    print("读取文档成功")
except PermissionError as e:
    print(f"读取被拒绝: {e}")

# 测试删除（应该被拒绝）
try:
    delete_document(
        resource_type='document',
        resource_id='doc-123',
        owner='team-a',
        session_id='sess-1'
    )
    print("删除文档成功")
except PermissionError as e:
    print(f"删除被拒绝: {e}")