feat(agent): RuntimeConfigMiddleware 添加自定义字段配置

This commit is contained in:
Wenjie Zhang 2026-01-30 12:40:54 +08:00
parent 9b742fd90a
commit 8cbebf4cca

View File

@ -18,27 +18,52 @@ class RuntimeConfigMiddleware(AgentMiddleware):
注意所有可能用到的知识库工具必须在初始化时预加载并注册到 self.tools 注意所有可能用到的知识库工具必须在初始化时预加载并注册到 self.tools
运行时根据配置从 self.tools 中筛选工具不能动态添加新工具 运行时根据配置从 self.tools 中筛选工具不能动态添加新工具
支持自定义上下文字段名称以便在不同场景如主智能体/子智能体使用不同的配置字段
""" """
def __init__(self, *, extra_tools: list[Any] | None = None): def __init__(
self,
*,
extra_tools: list[Any] | None = None,
model_context_name: str = "model",
system_prompt_context_name: str = "system_prompt",
tools_context_name: str = "tools",
knowledges_context_name: str = "knowledges",
mcps_context_name: str = "mcps",
):
"""初始化中间件 """初始化中间件
Args: Args:
extra_tools: 额外工具列表 create_agent tools 参数传入 extra_tools: 额外工具列表 create_agent tools 参数传入
model_context_name: 上下文中的模型字段名称默认 "model"
system_prompt_context_name: 上下文中的系统提示词字段名称默认 "system_prompt"
tools_context_name: 上下文中的工具列表字段名称默认 "tools"
knowledges_context_name: 上下文中的知识库列表字段名称默认 "knowledges"
mcps_context_name: 上下文中的 MCP 服务器列表字段名称默认 "mcps"
""" """
super().__init__() super().__init__()
# 这里的工具只是提供给 langchain 调用,并不是真正的绑定在模型上(后续会过滤)
self.kb_tools = get_kb_based_tools() self.kb_tools = get_kb_based_tools()
self.buildin_tools = get_buildin_tools() self.buildin_tools = get_buildin_tools()
self.tools = self.kb_tools + self.buildin_tools + (extra_tools or []) self.tools = self.kb_tools + self.buildin_tools + (extra_tools or [])
logger.debug(f"Initialized tools: {len(self.tools)}") # 存储自定义字段名称
self.model_context_name = model_context_name
self.system_prompt_context_name = system_prompt_context_name
self.tools_context_name = tools_context_name
self.knowledges_context_name = knowledges_context_name
self.mcps_context_name = mcps_context_name
logger.debug(
f"Initialized RuntimeConfigMiddleware with custom field names: model={model_context_name}, "
f"system_prompt={system_prompt_context_name}, tools={tools_context_name}, "
f"knowledges={knowledges_context_name}, mcps={mcps_context_name}"
)
async def awrap_model_call( async def awrap_model_call(
self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse] self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
) -> ModelResponse: ) -> ModelResponse:
runtime_context = request.runtime.context runtime_context = request.runtime.context
model = load_chat_model(getattr(runtime_context, "model", None)) model = load_chat_model(getattr(runtime_context, self.model_context_name, None))
enabled_tools = await self.get_tools_from_context(runtime_context) enabled_tools = await self.get_tools_from_context(runtime_context)
existing_tools = list(request.tools or []) existing_tools = list(request.tools or [])
@ -50,7 +75,7 @@ class RuntimeConfigMiddleware(AgentMiddleware):
# 动态生成 system message添加当前时间 # 动态生成 system message添加当前时间
cur_datetime = f"当前时间:{shanghai_now().strftime('%Y-%m-%d %H:%M:%S')} UTC" cur_datetime = f"当前时间:{shanghai_now().strftime('%Y-%m-%d %H:%M:%S')} UTC"
system_prompt = getattr(runtime_context, "system_prompt", "") or "" system_prompt = getattr(runtime_context, self.system_prompt_context_name, "") or ""
new_content = list(request.system_message.content_blocks) + [ new_content = list(request.system_message.content_blocks) + [
{"type": "text", "text": f"{cur_datetime}\n\n{system_prompt}"} {"type": "text", "text": f"{cur_datetime}\n\n{system_prompt}"}
] ]
@ -63,26 +88,28 @@ class RuntimeConfigMiddleware(AgentMiddleware):
async def get_tools_from_context(self, context) -> list: async def get_tools_from_context(self, context) -> list:
"""从上下文配置中获取工具列表""" """从上下文配置中获取工具列表"""
# 1. 基础工具 (从 context.tools 中筛选)
selected_tools = [] selected_tools = []
if context.tools: # 1. 基础工具 (从 context.tools 中筛选)
# 创建工具映射表 tools = getattr(context, self.tools_context_name, None)
if tools:
tools_map = {t.name: t for t in self.tools} tools_map = {t.name: t for t in self.tools}
for tool_name in context.tools: for tool_name in tools:
if tool_name in tools_map: if tool_name in tools_map:
selected_tools.append(tools_map[tool_name]) selected_tools.append(tools_map[tool_name])
else: else:
logger.warning(f"Tool '{tool_name}' not found in available tools. {tools_map.keys()=}") logger.warning(f"Tool '{tool_name}' not found in available tools. {tools_map.keys()=}")
# 2. 知识库工具 # 2. 知识库工具
if context.knowledges: knowledges = getattr(context, self.knowledges_context_name, None)
kb_tools = get_kb_based_tools(db_names=context.knowledges) if knowledges:
kb_tools = get_kb_based_tools(db_names=knowledges)
selected_tools.extend(kb_tools) selected_tools.extend(kb_tools)
# 3. MCP 工具(使用统一入口,自动过滤 disabled_tools # 3. MCP 工具(使用统一入口,自动过滤 disabled_tools
if context.mcps: mcps = getattr(context, self.mcps_context_name, None)
for server_name in context.mcps: if mcps:
for server_name in mcps:
mcp_tools = await get_enabled_mcp_tools(server_name) mcp_tools = await get_enabled_mcp_tools(server_name)
selected_tools.extend(mcp_tools) selected_tools.extend(mcp_tools)