From a9b5e84618bffa296e6357f3eb113e4ff8a2f65b Mon Sep 17 00:00:00 2001
From: Wenjie Zhang <xerrors@163.com>
Date: Fri, 27 Jun 2025 01:49:31 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=E5=BC=82=E6=AD=A5?=
 =?UTF-8?q?=E5=B7=A5=E5=85=B7=E6=94=AF=E6=8C=81=EF=BC=8C=E4=BC=98=E5=8C=96?=
 =?UTF-8?q?=E5=B7=A5=E5=85=B7=E6=A3=80=E7=B4=A2=E5=92=8C=20LLM=20=E8=B0=83?=
 =?UTF-8?q?=E7=94=A8=E9=80=BB=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 在 tools_factory.py 中引入异步检索器包装函数，确保异步工具的正确处理。
- 更新 ChatbotAgent 类的 llm_call 方法为异步版本，以支持异步工具调用。
---
 src/agents/chatbot/graph.py |  7 ++++---
 src/agents/tools_factory.py | 25 ++++++++++++++++++++++---
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/src/agents/chatbot/graph.py b/src/agents/chatbot/graph.py
index 902a4eea..d0091b6a 100644
--- a/src/agents/chatbot/graph.py
+++ b/src/agents/chatbot/graph.py
@@ -45,8 +45,8 @@ class ChatbotAgent(BaseAgent):
             logger.info(f"使用工具: {tool_names}")
             return [platform_tools[tool] for tool in tool_names]
 
-    def llm_call(self, state: State, config: RunnableConfig = None) -> dict[str, Any]:
-        """调用 llm 模型"""
+    async def llm_call(self, state: State, config: RunnableConfig = None) -> dict[str, Any]:
+        """调用 llm 模型 - 异步版本以支持异步工具"""
         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()}"
@@ -55,7 +55,8 @@ class ChatbotAgent(BaseAgent):
         if tools := self._get_tools(conf.tools):
             model = model.bind_tools(tools)
 
-        res = model.invoke(
+        # 使用异步调用
+        res = await model.ainvoke(
             [{"role": "system", "content": system_prompt}, *state["messages"]]
         )
         return {"messages": [res]}
diff --git a/src/agents/tools_factory.py b/src/agents/tools_factory.py
index 69c4461c..feb75d03 100644
--- a/src/agents/tools_factory.py
+++ b/src/agents/tools_factory.py
@@ -2,6 +2,8 @@ import json
 import re
 from collections.abc import Callable
 from typing import Annotated, Any
+import asyncio
+import logging
 
 from langchain_tavily import TavilySearch
 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():
-        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 = (
             f"使用 {retrieve_info['name']} 知识库进行检索。\n"
             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(
-            retrieve_info["retriever"],
+            coroutine=async_retriever_wrapper,  # 指定为协程
             name=name,
             description=description,
-            args_schema=KnowledgeRetrieverModel)
+            args_schema=KnowledgeRetrieverModel
+        )
 
     return tools