feat: 添加异步工具支持,优化工具检索和 LLM 调用逻辑
- 在 tools_factory.py 中引入异步检索器包装函数,确保异步工具的正确处理。 - 更新 ChatbotAgent 类的 llm_call 方法为异步版本,以支持异步工具调用。
This commit is contained in:
parent
be3adf7b9d
commit
a9b5e84618
@ -45,8 +45,8 @@ class ChatbotAgent(BaseAgent):
|
|||||||
logger.info(f"使用工具: {tool_names}")
|
logger.info(f"使用工具: {tool_names}")
|
||||||
return [platform_tools[tool] for tool in tool_names]
|
return [platform_tools[tool] for tool in tool_names]
|
||||||
|
|
||||||
def llm_call(self, state: State, config: RunnableConfig = None) -> dict[str, Any]:
|
async def llm_call(self, state: State, config: RunnableConfig = None) -> dict[str, Any]:
|
||||||
"""调用 llm 模型"""
|
"""调用 llm 模型 - 异步版本以支持异步工具"""
|
||||||
conf = self.config_schema.from_runnable_config(config, agent_name=self.name)
|
conf = self.config_schema.from_runnable_config(config, agent_name=self.name)
|
||||||
|
|
||||||
system_prompt = f"{conf.system_prompt} Now is {get_cur_time_with_utc()}"
|
system_prompt = f"{conf.system_prompt} Now is {get_cur_time_with_utc()}"
|
||||||
@ -55,7 +55,8 @@ class ChatbotAgent(BaseAgent):
|
|||||||
if tools := self._get_tools(conf.tools):
|
if tools := self._get_tools(conf.tools):
|
||||||
model = model.bind_tools(tools)
|
model = model.bind_tools(tools)
|
||||||
|
|
||||||
res = model.invoke(
|
# 使用异步调用
|
||||||
|
res = await model.ainvoke(
|
||||||
[{"role": "system", "content": system_prompt}, *state["messages"]]
|
[{"role": "system", "content": system_prompt}, *state["messages"]]
|
||||||
)
|
)
|
||||||
return {"messages": [res]}
|
return {"messages": [res]}
|
||||||
|
|||||||
@ -2,6 +2,8 @@ import json
|
|||||||
import re
|
import re
|
||||||
from collections.abc import Callable
|
from collections.abc import Callable
|
||||||
from typing import Annotated, Any
|
from typing import Annotated, Any
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
|
||||||
from langchain_tavily import TavilySearch
|
from langchain_tavily import TavilySearch
|
||||||
from langchain_core.tools import BaseTool, StructuredTool, tool
|
from langchain_core.tools import BaseTool, StructuredTool, tool
|
||||||
@ -81,16 +83,33 @@ def get_all_tools():
|
|||||||
|
|
||||||
# 获取所有知识库
|
# 获取所有知识库
|
||||||
for db_Id, retrieve_info in knowledge_base.get_retrievers().items():
|
for db_Id, retrieve_info in knowledge_base.get_retrievers().items():
|
||||||
name = f"retrieve_{db_Id}" # Deepseek does not support non-alphanumeric characters in tool names
|
name = f"retrieve_{db_Id[:8]}" # Deepseek does not support non-alphanumeric characters in tool names
|
||||||
description = (
|
description = (
|
||||||
f"使用 {retrieve_info['name']} 知识库进行检索。\n"
|
f"使用 {retrieve_info['name']} 知识库进行检索。\n"
|
||||||
f"下面是这个知识库的描述:\n{retrieve_info['description']}"
|
f"下面是这个知识库的描述:\n{retrieve_info['description']}"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# 创建异步工具,确保正确处理异步检索器
|
||||||
|
async def async_retriever_wrapper(query_text: str, db_id=db_Id):
|
||||||
|
"""异步检索器包装函数"""
|
||||||
|
retriever = retrieve_info["retriever"]
|
||||||
|
try:
|
||||||
|
if asyncio.iscoroutinefunction(retriever):
|
||||||
|
result = await retriever(query_text)
|
||||||
|
else:
|
||||||
|
result = retriever(query_text)
|
||||||
|
return result
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Error in retriever {db_id}: {e}")
|
||||||
|
return f"检索失败: {str(e)}"
|
||||||
|
|
||||||
|
# 使用 StructuredTool.from_function 创建异步工具
|
||||||
tools[name] = StructuredTool.from_function(
|
tools[name] = StructuredTool.from_function(
|
||||||
retrieve_info["retriever"],
|
coroutine=async_retriever_wrapper, # 指定为协程
|
||||||
name=name,
|
name=name,
|
||||||
description=description,
|
description=description,
|
||||||
args_schema=KnowledgeRetrieverModel)
|
args_schema=KnowledgeRetrieverModel
|
||||||
|
)
|
||||||
|
|
||||||
return tools
|
return tools
|
||||||
|
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user