1. AI Agent开发实战:从工具调用到安全防护
最近在开发一个能处理公司内部数据的AI Agent时,踩了不少坑,也积累了一些实战经验。这个Agent不仅能查询公司机密文档,还能进行精确计算,比普通聊天机器人强大得多。下面我就把整个开发过程、技术实现和安全防护方案详细分享给大家。
2. AI Agent核心架构解析
2.1 什么是真正的AI Agent?
很多人以为AI Agent就是高级版的ChatGPT,其实远不止如此。一个完整的AI Agent应该具备四大核心能力:
- LLM核心:使用大语言模型进行推理和决策
- 记忆系统:包括短期记忆(对话历史)和长期记忆(RAG知识库)
- 规划能力:能拆解复杂任务并制定执行流程
- 工具调用:可以自主调用外部函数和API
我开发的这个Agent主要聚焦在工具调用能力上,使用了LangChain框架来实现。
2.2 工具调用的技术选型
在工具调用方案上,我对比了几种主流方案:
| 方案 | 优点 | 缺点 | 适用场景 |
|---|---|---|---|
| LangChain工具调用 | 开发简单,生态完善 | 性能开销较大 | 快速原型开发 |
| 直接API调用 | 性能高 | 开发复杂 | 生产环境 |
| 自定义插件 | 灵活性高 | 维护成本高 | 特殊需求 |
最终选择LangChain是因为:
- 它提供了完整的工具调用流程封装
- 内置了安全防护机制
- 社区支持好,遇到问题容易找到解决方案
3. 实战开发:构建多功能Agent
3.1 环境准备与依赖安装
首先需要安装必要的Python包:
bash复制pip install langchain langchain-community faiss-cpu
这里特别说明几个关键依赖:
langchain: 提供Agent开发的核心框架faiss-cpu: 用于本地向量存储和检索dashscope: 通义千问的Python SDK
注意:如果使用GPU环境,可以安装faiss-gpu版本来加速向量检索。
3.2 核心工具函数实现
我实现了两个核心工具函数:
3.2.1 精确计算器
python复制@tool
def calculator(expression: str) -> str:
"""
计算数学表达式。需要精确计算时使用。
参数:
expression: 数学算式,如 "2 + 2" 或 "500 * 0.8"。
返回:
str: 计算结果,如 "4.0" 或 "400.0"。
"""
print(f" [工具调用] 计算器正在计算: {expression}")
try:
return str(eval(expression))
except Exception as e:
return f"计算错误: {e}"
3.2.2 RAG文档检索
python复制@tool
def rag_search(query: str) -> str:
"""
从数据库中搜索与查询公司内部相关的文档...
"""
# 文档预处理和向量化存储
raw_text = """【公司内部机密:代号"深蓝计划"...】"""
RAG_PATH = "faiss_index"
# 文档切分和向量化
docs = [Document(page_content=raw_text)]
text_splitter = RecursiveCharacterTextSplitter(chunk_size=25, chunk_overlap=5)
split_docs = text_splitter.split_documents(docs)
embeddings = DashScopeEmbeddings(model="text-embedding-v1")
# 加载或创建向量数据库
if os.path.exists(RAG_PATH):
ragdb = FAISS.load_local(RAG_PATH, embeddings, allow_dangerous_deserialization=True)
else:
ragdb = FAISS.from_documents(split_docs, embeddings)
ragdb.save_local(RAG_PATH)
return "\n\n".join(doc.page_content for doc in ragdb.similarity_search(query, k=2))
3.3 Agent核心执行流程
Agent的执行采用了多轮对话机制,关键代码如下:
python复制def run_agent(query:str):
# 初始化工具和模型
tool_maps = {"rag_search": rag_search, "calculator": calculator}
llm = ChatTongyi(model_name="qwen-plus")
tool_llm = llm.bind_tools(tools=list(tool_maps.values()))
message = [HumanMessage(content=query)]
for i in range(5): # 最多5轮对话防止死循环
response = tool_llm.invoke(message)
message.append(response)
if not response.tool_calls: # 无需工具调用
return response.content
# 处理工具调用
for tool_call in response.tool_calls:
func_name = tool_call["name"]
if func_name in tool_maps: # 安全检查
tool_output = tool_maps[func_name].invoke(tool_call["args"])
message.append(ToolMessage(content=tool_output, tool_call_id=tool_call["id"]))
这个流程实现了完整的工具调用闭环:
- 用户输入问题
- LLM判断是否需要调用工具
- 执行工具并获取结果
- 将结果返回给LLM生成最终回复
4. 安全防护实战方案
4.1 识别安全风险
在开发过程中,我发现了几处潜在的安全隐患:
- eval注入风险:计算器函数直接使用eval(),可能被恶意利用
- 工具滥用风险:Agent可能被诱导调用不该调用的工具
- 信息泄露风险:RAG检索可能返回敏感信息
4.2 安全加固措施
4.2.1 计算器安全改造
原始版本直接使用eval非常危险,我做了以下改进:
python复制import ast
import operator
def safe_eval(expr: str) -> str:
# 允许的操作符白名单
ALLOWED_OPERATORS = {
ast.Add: operator.add,
ast.Sub: operator.sub,
ast.Mult: operator.mul,
ast.Div: operator.truediv,
ast.Pow: operator.pow,
ast.USub: operator.neg
}
# 允许的节点类型
ALLOWED_NODES = (ast.Expression, ast.Num, ast.BinOp, ast.UnaryOp)
try:
tree = ast.parse(expr, mode='eval')
for node in ast.walk(tree):
if not isinstance(node, ALLOWED_NODES):
raise ValueError(f"禁止的操作: {type(node).__name__}")
return str(eval(compile(tree, '', 'eval'), {}, ALLOWED_OPERATORS))
except Exception as e:
return f"计算错误: {e}"
这个安全版本:
- 使用AST解析表达式
- 严格限制允许的操作符和节点类型
- 完全隔离执行环境
4.2.2 工具调用防护
在工具调用环节增加了多重防护:
- 工具白名单检查:
python复制if func_name not in tool_maps:
return "错误: 未授权的工具调用"
- 输入参数过滤:
python复制# 对字符串参数进行HTML转义
from html import escape
safe_args = {k: escape(str(v)) for k,v in tool_call["args"].items()}
- 调用频率限制:
python复制from datetime import datetime, timedelta
class ToolRateLimiter:
def __init__(self, calls=5, per_seconds=60):
self.calls = calls
self.per_seconds = per_seconds
self.history = []
def check(self):
now = datetime.now()
self.history = [t for t in self.history if now - t < timedelta(seconds=self.per_seconds)]
if len(self.history) >= self.calls:
return False
self.history.append(now)
return True
4.3 隐私数据保护
对于RAG系统中的敏感数据,我采取了以下措施:
- 数据脱敏处理:
python复制def anonymize(text: str) -> str:
# 替换敏感信息
patterns = {
r"\d{4}-\d{2}-\d{2}": "[DATE]",
r"\d+元人民币": "[AMOUNT]",
r"代号\"\w+\"": "[PROJECT_CODE]"
}
for pat, repl in patterns.items():
text = re.sub(pat, repl, text)
return text
- 访问控制:
python复制def check_access(user: str, doc_id: str) -> bool:
# 实现基于角色的访问控制
user_roles = get_user_roles(user)
doc_acl = get_document_acl(doc_id)
return bool(set(user_roles) & set(doc_acl))
5. 实战经验与避坑指南
5.1 调试技巧
在开发过程中,我发现这些调试方法特别有用:
- 完整日志记录:
python复制import logging
logging.basicConfig(
level=logging.DEBUG,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[logging.FileHandler('agent.log'), logging.StreamHandler()]
)
- 中间结果检查:
python复制# 在关键步骤插入检查点
def debug_step(message):
print(f"=== DEBUG ===\n{json.dumps(message, indent=2, ensure_ascii=False)}")
return message
5.2 性能优化
经过测试,我发现这些优化措施能显著提升Agent性能:
- 向量检索优化:
python复制# 使用更高效的切分策略
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=150, # 适当增大chunk大小
chunk_overlap=30,
length_function=len,
is_separator_regex=False
)
- LLM调用缓存:
python复制from langchain.cache import InMemoryCache
from langchain.globals import set_llm_cache
set_llm_cache(InMemoryCache()) # 简单的内存缓存
5.3 常见问题解决
在实际使用中,我遇到了这些问题和解决方案:
- 工具调用死循环
- 现象:Agent不断要求调用同一个工具
- 解决:严格限制最大轮次(如5轮),并添加循环检测逻辑
- RAG检索不准
- 现象:返回不相关的文档片段
- 解决:调整chunk大小和重叠比例,优化embedding模型
- 计算精度问题
- 现象:浮点数计算出现精度误差
- 解决:使用decimal模块进行高精度计算
python复制from decimal import Decimal, getcontext
getcontext().prec = 10 # 设置计算精度
def precise_calc(expr: str) -> str:
try:
return str(eval(expr, {'__builtins__': None}, {'Decimal': Decimal}))
except Exception as e:
return f"计算错误: {e}"
开发AI Agent既充满挑战又极具价值。通过这次项目,我深刻体会到安全防护的重要性。建议大家在实际开发中,一定要把安全考虑放在首位,从设计阶段就建立完善的安全防护机制。