构建安全的“防具”—— 在AI项目中用ADT杜绝数据污染
问题场景裸露数据结构的风险在原系统中消息列表以裸露的Python列表形式在各个函数间传递# 原代码裸露的数据结构 final_messages [] final_messages.append({role: system, content: 你是一个助手}) # 某处代码不小心修改了数据 some_messages final_messages some_messages.append({role: hacker, content: 恶意消息}) # 污染了原始数据 # 另一处代码也难以保证不修改 final_messages[0][content] 被篡改的内容 # 直接修改这种设计存在几个严重问题数据可以被任意修改没有保护机制消息格式没有验证role 是否合法是否包含 content代码意图不清晰难以追踪数据流向ADT 设计SessionContext类定义与表示不变量# app/domain/session_context.py class SessionContext: ADT 代表一次智能体对话的完整短期记忆上下文。 AF: 代表一个对话会话的状态包含历史消息列表和元数据。 RI (表示不变量): 1. _messages 列表不为 None且每个元素都是字典 2. 每条消息必须包含 role 和 content 字段 3. role 只能是 user, assistant, system 之一 4. _messages 长度不超过 100 条 ALLOWED_ROLES {user, assistant, system} MAX_MESSAGES 100 def __init__(self, messages: Optional[List[Dict]] None): self._messages copy.deepcopy(messages) if messages else [] self._check_rep() # 确保初始状态合法防御性编程设值与取值都防守def add_message(self, message: Dict[str, str]) - None: 添加消息 - 设值时防守 # 1. 深拷贝断开外部引用 new_message copy.deepcopy(message) # 2. 验证消息格式 if role not in new_message or content not in new_message: raise ValueError(消息必须包含 role 和 content 字段) if new_message[role] not in self.ALLOWED_ROLES: raise ValueError(frole 必须是 {self.ALLOWED_ROLES} 之一) if len(self._messages) self.MAX_MESSAGES: raise ValueError(f消息数量已达上限 {self.MAX_MESSAGES}) # 3. 添加消息 self._messages.append(new_message) # 4. 验证不变量 self._check_rep() def get_messages(self) - List[Dict[str, str]]: 获取消息 - 取值时防守 # 关键返回深拷贝而不是内部引用 return copy.deepcopy(self._messages)RI 检查时刻保证对象合法def _check_rep(self) - None: 检查表示不变量 assert self._messages is not None, RI 违反: _messages 为 None for i, msg in enumerate(self._messages): assert isinstance(msg, dict), f消息 {i} 不是字典 assert role in msg, f消息 {i} 缺少 role assert content in msg, f消息 {i} 缺少 content assert msg[role] in self.ALLOWED_ROLES, \ f消息 {i} 的 role {msg[role]} 不在允许列表中 assert len(self._messages) self.MAX_MESSAGES, \ f消息数量 {len(self._messages)} 超过上限破坏性测试证明防具有效# tests/domain/test_session_context.py def test_add_invalid_role(): 拦截非法角色抛出 ValueError ctx SessionContext() with pytest.raises(ValueError, matchrole 必须是): ctx.add_message({role: invalid_role, content: test}) def test_get_messages_returns_copy(): 外部修改副本不影响内部 ctx SessionContext() ctx.add_message({role: user, content: 原始消息}) messages_copy ctx.get_messages() messages_copy.append({role: assistant, content: 恶意添加}) messages_copy[0][content] 篡改 # 内部状态不受影响 assert ctx.get_message_count() 1 assert ctx.get_messages()[0][content] 原始消息 def test_add_message_deep_copy(): 外部修改传入的字典不影响内部 ctx SessionContext() external_msg {role: user, content: 外部消息} ctx.add_message(external_msg) # 外部修改原字典 external_msg[content] 被篡改了 external_msg[role] hacker # 内部消息未被修改 internal_msg ctx.get_messages()[0] assert internal_msg[content] 外部消息 assert internal_msg[role] user进阶不可变 KnowledgeChunk除了消息列表知识库切块也是数据污染的重灾区。设计了不可变的 KnowledgeChunk 对象# app/domain/knowledge_chunk.py from dataclasses import dataclass, field import copy dataclass(frozenTrue) # frozenTrue 使对象不可变 class KnowledgeChunk: 知识库文档切块的不可变对象。 AF: 代表一个文档片段包含文本、元数据和唯一标识。 text: str chunk_id: str metadata: Dict[str, Any] field(default_factorydict) embedding: Optional[list] field(defaultNone, compareFalse) def __post_init__(self): 验证不可变对象的状态 if not isinstance(self.text, str) or not self.text.strip(): raise ValueError(text must be a non-empty string) # 冻结 metadata 的深层内容 object.__setattr__(self, metadata, copy.deepcopy(self.metadata))使用效果# 重构前裸露字典 chunk {text: 内容, metadata: {file_id: 123}} chunk[text] 被篡改 # 可以直接修改 chunk[new_key] 任意添加 # 重构后不可变对象 chunk KnowledgeChunk(text内容, chunk_idc1, metadata{file_id: 123}) chunk.text 被篡改 # ❌ dataclass(frozenTrue) 禁止修改总结核心收获设计要点具体实践表示不变量 (RI)明确定义合法状态用 _check_rep() 强制保证防御性复制输入时 deepcopy输出时也 deepcopy不可变对象使用 dataclass(frozenTrue) 从根源杜绝修改表示独立性内部用私有变量外部只能通过方法访问ADT 的价值数据安全外部无法意外修改内部状态代码可读ADT 的方法名清晰表达意图提前容错非法数据在进入时就报错而非运行时崩溃易于测试ADT 的接口明确可以独立编写单元测试