添加自定义工具功能，给未来的 MCP 打一下基础
This commit is contained in:
Wenjie Zhang
parent
0bd862069c
commit
693d05d808
@ -28,7 +28,7 @@ async def update_config(key = Body(...), value = Body(...)):
|
||||
@base.post("/restart")
|
||||
async def restart():
|
||||
knowledge_base.restart()
|
||||
graph_base.restart()
|
||||
graph_base.start()
|
||||
retriever.restart()
|
||||
return {"message": "Restarted!"}
|
||||
|
||||
|
||||
@ -12,6 +12,7 @@ from src.core import HistoryManager
|
||||
from src.agents import agent_manager
|
||||
from src.models import select_model
|
||||
from src.utils.logging_config import logger
|
||||
from src.agents.tools_factory import get_all_tools
|
||||
|
||||
chat = APIRouter(prefix="/chat")
|
||||
|
||||
@ -201,3 +202,8 @@ async def get_chat_models(model_provider: str):
|
||||
"""获取指定模型提供商的模型列表"""
|
||||
model = select_model(model_provider=model_provider)
|
||||
return {"models": model.get_models()}
|
||||
|
||||
@chat.get("/tools")
|
||||
async def get_tools():
|
||||
"""获取所有工具"""
|
||||
return {"tools": list(get_all_tools().keys())}
|
||||
|
||||
@ -34,3 +34,12 @@ class ChatbotConfiguration(Configuration):
|
||||
"description": "智能体的驱动模型"
|
||||
},
|
||||
)
|
||||
|
||||
tools: list[str] = field(
|
||||
default_factory=list,
|
||||
metadata={
|
||||
"name": "工具",
|
||||
"configurable": False,
|
||||
"description": "工具列表"
|
||||
},
|
||||
)
|
||||
|
||||
@ -13,31 +13,32 @@ from src.utils import logger
|
||||
from src.agents.registry import State, BaseAgent
|
||||
from src.agents.utils import load_chat_model, get_cur_time_with_utc
|
||||
from src.agents.chatbot.configuration import ChatbotConfiguration
|
||||
from src.agents.tools_factory import _TOOLS_REGISTRY
|
||||
from src.agents.tools_factory import get_all_tools
|
||||
|
||||
class ChatbotAgent(BaseAgent):
|
||||
name = "chatbot"
|
||||
description = "A chatbot that can answer questions and help with tasks."
|
||||
description = "基础的对话机器人，可以回答问题，默认不使用任何工具，可在配置中启用需要的工具。"
|
||||
requirements = ["TAVILY_API_KEY", "ZHIPUAI_API_KEY"]
|
||||
all_tools = ["TavilySearchResults", "multiply", "add", "subtract", "divide"]
|
||||
config_schema = ChatbotConfiguration
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def _get_tools(self, config_schema: RunnableConfig):
|
||||
"""根据配置获取工具，如果配置为空，则使用所有工具，如果配置为列表，则使用列表中的工具，
|
||||
如果配置为其他类型，则抛出错误"""
|
||||
conf_tools = config_schema.get("tools")
|
||||
if conf_tools == None:
|
||||
tool_names = self.all_tools
|
||||
elif isinstance(conf_tools, list):
|
||||
tool_names = [tool for tool in self.all_tools if tool in conf_tools]
|
||||
def _get_tools(self, tools: list[str]):
|
||||
"""根据配置获取工具。
|
||||
默认不使用任何工具。
|
||||
如果配置为列表，则使用列表中的工具。
|
||||
"""
|
||||
platform_tools = get_all_tools()
|
||||
if tools is None or not isinstance(tools, list) or len(tools) == 0:
|
||||
# 默认不使用任何工具
|
||||
logger.info("未配置工具或配置为空，不使用任何工具")
|
||||
return []
|
||||
else:
|
||||
raise ValueError(f"tools 配置错误: {conf_tools}")
|
||||
|
||||
logger.info(f"Tools: {tool_names}")
|
||||
return [_TOOLS_REGISTRY[tool] for tool in tool_names]
|
||||
# 使用配置中指定的工具
|
||||
tool_names = [tool for tool in platform_tools.keys() if tool in tools]
|
||||
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 模型"""
|
||||
@ -46,7 +47,7 @@ class ChatbotAgent(BaseAgent):
|
||||
|
||||
system_prompt = f"{conf.system_prompt} Now is {get_cur_time_with_utc()}"
|
||||
model = load_chat_model(conf.model)
|
||||
model_with_tools = model.bind_tools(self._get_tools(config_schema))
|
||||
model_with_tools = model.bind_tools(self._get_tools(conf.tools))
|
||||
logger.info(f"llm_call with config: {conf}, {conf.model}")
|
||||
|
||||
res = model_with_tools.invoke(
|
||||
@ -56,9 +57,10 @@ class ChatbotAgent(BaseAgent):
|
||||
|
||||
def get_graph(self, config_schema: RunnableConfig = None, **kwargs):
|
||||
"""构建图"""
|
||||
conf = self.config_schema.from_runnable_config(config_schema)
|
||||
workflow = StateGraph(State, config_schema=self.config_schema)
|
||||
workflow.add_node("chatbot", self.llm_call)
|
||||
workflow.add_node("tools", ToolNode(tools=self._get_tools(config_schema)))
|
||||
workflow.add_node("tools", ToolNode(tools=self._get_tools(conf.tools)))
|
||||
workflow.add_edge(START, "chatbot")
|
||||
workflow.add_conditional_edges(
|
||||
"chatbot",
|
||||
|
||||
@ -2,6 +2,8 @@ import os
|
||||
|
||||
from langchain_openai import ChatOpenAI
|
||||
|
||||
from src import graph_base
|
||||
|
||||
model = ChatOpenAI(model="glm-4-plus",
|
||||
api_key=os.getenv("ZHIPUAI_API_KEY"),
|
||||
base_url="https://open.bigmodel.cn/api/paas/v4/",
|
||||
@ -10,26 +12,14 @@ model = ChatOpenAI(model="glm-4-plus",
|
||||
|
||||
# For this tutorial we will use custom tool that returns pre-defined values for weather in two cities (NYC & SF)
|
||||
|
||||
from typing import Literal
|
||||
from typing import Literal, Annotated
|
||||
|
||||
from langchain_core.tools import tool
|
||||
from langchain_core.tools import tool, StructuredTool
|
||||
|
||||
|
||||
@tool
|
||||
def get_weather(city: Literal["nyc", "sf"]):
|
||||
"""Use this to get weather information."""
|
||||
if city == "nyc":
|
||||
return "It might be cloudy in nyc"
|
||||
elif city == "sf":
|
||||
return "It's always sunny in sf"
|
||||
else:
|
||||
raise AssertionError("Unknown city")
|
||||
|
||||
tools = []
|
||||
|
||||
tools = [get_weather]
|
||||
|
||||
|
||||
# Define the graph
|
||||
|
||||
from langgraph.prebuilt import create_react_agent
|
||||
|
||||
|
||||
@ -1,13 +1,15 @@
|
||||
import json
|
||||
import re
|
||||
import os
|
||||
from typing import Any, Callable, Optional, Type, Union
|
||||
from typing import Any, Callable, Optional, Type, Union, Annotated
|
||||
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from langchain_core.tools import tool, BaseTool
|
||||
from langchain_core.tools import tool, BaseTool, StructuredTool
|
||||
from langchain_community.tools.tavily_search import TavilySearchResults
|
||||
|
||||
from src import graph_base, knowledge_base
|
||||
|
||||
# refs https://github.com/chatchat-space/LangGraph-Chatchat chatchat-server/chatchat/server/agent/tools_factory/tools_registry.py
|
||||
def regist_tool(
|
||||
*args: Any,
|
||||
@ -63,6 +65,24 @@ def regist_tool(
|
||||
return t
|
||||
|
||||
|
||||
class KnowledgeRetrieverModel(BaseModel):
|
||||
query: str = Field(description="The query to get knowledge graph.")
|
||||
|
||||
|
||||
|
||||
def get_all_tools():
|
||||
"""获取所有工具"""
|
||||
tools = _TOOLS_REGISTRY.copy()
|
||||
for db_Id, retrieve_info in knowledge_base.get_retrievers().items():
|
||||
name = f"retrieve_{retrieve_info['name']}"
|
||||
tools[name] = StructuredTool.from_function(
|
||||
retrieve_info["retriever"],
|
||||
name=name,
|
||||
description=retrieve_info["description"],
|
||||
args_schema=KnowledgeRetrieverModel)
|
||||
|
||||
return tools
|
||||
|
||||
class BaseToolOutput:
|
||||
"""
|
||||
LLM 要求 Tool 的输出为 str，但 Tool 用在别处时希望它正常返回结构化数据。
|
||||
@ -92,33 +112,30 @@ class BaseToolOutput:
|
||||
else:
|
||||
return str(self.data)
|
||||
|
||||
|
||||
@tool
|
||||
def calculator(a: float, b: float, operation: str) -> float:
|
||||
"""Calculate two numbers."""
|
||||
if operation == "add":
|
||||
return a + b
|
||||
elif operation == "subtract":
|
||||
return a - b
|
||||
elif operation == "multiply":
|
||||
return a * b
|
||||
elif operation == "divide":
|
||||
return a / b
|
||||
else:
|
||||
raise ValueError(f"Invalid operation: {operation}, only support add, subtract, multiply, divide")
|
||||
|
||||
@tool
|
||||
def multiply(first_int: int, second_int: int) -> int:
|
||||
"""Multiply two integers together."""
|
||||
return first_int * second_int
|
||||
def get_knowledge_graph(query: Annotated[str, "The query to get knowledge graph."]):
|
||||
"""Use this to get knowledge graph."""
|
||||
return graph_base.query_node(query, hops=2)
|
||||
|
||||
@tool
|
||||
def add(first_int: int, second_int: int) -> int:
|
||||
"""Add two integers together."""
|
||||
return first_int + second_int
|
||||
|
||||
@tool
|
||||
def subtract(first_int: int, second_int: int) -> int:
|
||||
"""Subtract two integers."""
|
||||
return first_int - second_int
|
||||
|
||||
@tool
|
||||
def divide(first_int: int, second_int: int) -> int:
|
||||
"""Divide two integers."""
|
||||
return first_int / second_int
|
||||
|
||||
|
||||
_TOOLS_REGISTRY = {
|
||||
"multiply": multiply,
|
||||
"add": add,
|
||||
"subtract": subtract,
|
||||
"divide": divide,
|
||||
"calculator": calculator,
|
||||
"TavilySearchResults": TavilySearchResults(max_results=10),
|
||||
"get_knowledge_graph": get_knowledge_graph,
|
||||
}
|
||||
|
||||
@ -209,6 +209,9 @@ class GraphDatabase:
|
||||
|
||||
def query_node(self, entity_name, hops=2, **kwargs):
|
||||
# TODO 添加判断节点数量为 0 停止检索
|
||||
# 判断是否启动
|
||||
if not self.is_running():
|
||||
raise Exception("图数据库未启动")
|
||||
|
||||
logger.debug(f"Query graph node {entity_name} with {hops=}")
|
||||
if kwargs.get("exact_match"):
|
||||
@ -266,7 +269,7 @@ class GraphDatabase:
|
||||
def query(tx, entity_name, hops):
|
||||
result = tx.run(f"""
|
||||
MATCH (n {{name: $entity_name}})-[r*1..{hops}]-(m)
|
||||
RETURN n, r, m
|
||||
RETURN n {{.*, embedding: null}} AS n, r, m {{.*, embedding: null}} AS m
|
||||
""", entity_name=entity_name)
|
||||
return result.values()
|
||||
|
||||
@ -279,7 +282,7 @@ class GraphDatabase:
|
||||
def query(tx, hops):
|
||||
result = tx.run(f"""
|
||||
MATCH (n)-[r*1..{hops}]->(m)
|
||||
RETURN n, r, m
|
||||
RETURN n {{.*, embedding: null}} AS n, r, m {{.*, embedding: null}} AS m
|
||||
""")
|
||||
return result.values()
|
||||
|
||||
@ -292,7 +295,7 @@ class GraphDatabase:
|
||||
def query(tx, relationship_type, hops):
|
||||
result = tx.run(f"""
|
||||
MATCH (n)-[r:`{relationship_type}`*1..{hops}]->(m)
|
||||
RETURN n, r, m
|
||||
RETURN n {{.*, embedding: null}} AS n, r, m {{.*, embedding: null}} AS m
|
||||
""")
|
||||
return result.values()
|
||||
|
||||
@ -307,7 +310,7 @@ class GraphDatabase:
|
||||
MATCH (n:Entity)
|
||||
WHERE n.name CONTAINS $keyword
|
||||
MATCH (n)-[r*1..{hops}]->(m)
|
||||
RETURN n, r, m
|
||||
RETURN n {{.*, embedding: null}} AS n, r, m {{.*, embedding: null}} AS m
|
||||
""", keyword=keyword)
|
||||
return result.values()
|
||||
|
||||
@ -321,7 +324,7 @@ class GraphDatabase:
|
||||
result = tx.run(f"""
|
||||
MATCH (n {{name: $node_name}})
|
||||
OPTIONAL MATCH (n)-[r*1..{hops}]->(m)
|
||||
RETURN n, r, m
|
||||
RETURN n {{.*, embedding: null}} AS n, r, m {{.*, embedding: null}} AS m
|
||||
""", node_name=node_name)
|
||||
return result.values()
|
||||
|
||||
|
||||
@ -344,7 +344,7 @@ class KnowledgeBase:
|
||||
"all_results": all_db_result,
|
||||
}
|
||||
|
||||
def get_retriever(self, db_id):
|
||||
def get_retriever_by_db_id(self, db_id):
|
||||
retriever_params = {
|
||||
"distance_threshold": self.default_distance_threshold,
|
||||
"rerank_threshold": self.default_rerank_threshold,
|
||||
@ -358,6 +358,16 @@ class KnowledgeBase:
|
||||
|
||||
return retriever
|
||||
|
||||
def get_retrievers(self):
|
||||
retrievers = {}
|
||||
for db in self.db_manager.get_all_databases():
|
||||
retrievers[db["db_id"]] = {
|
||||
"name": db["name"],
|
||||
"description": db["description"],
|
||||
"retriever": self.get_retriever_by_db_id(db["db_id"]),
|
||||
}
|
||||
return retrievers
|
||||
|
||||
################################
|
||||
#* Below is the code for milvus #
|
||||
################################
|
||||
|
||||
@ -360,6 +360,8 @@ const sendMessageWithText = async (text) => {
|
||||
body: JSON.stringify(requestData)
|
||||
});
|
||||
|
||||
console.log("requestData", requestData);
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error('请求失败');
|
||||
}
|
||||
@ -554,7 +556,7 @@ const handleFinished = async (data) => {
|
||||
const handleMessageById = async (data) => {
|
||||
const msgId = data.msg.id;
|
||||
const msgType = data.msg.type;
|
||||
console.log("data", data);
|
||||
// console.log("data", data);
|
||||
|
||||
// 查找现有消息
|
||||
const existingMsgIndex = messageMap.value.get(msgId);
|
||||
|
||||
@ -163,6 +163,23 @@
|
||||
</a-form-item>
|
||||
</template>
|
||||
|
||||
<!-- 添加工具选择部分 -->
|
||||
<a-form-item label="可用工具" name="tools" class="config-item">
|
||||
<p class="description">选择要启用的工具</p>
|
||||
<a-form-item-rest>
|
||||
<div class="tools-switches">
|
||||
<div v-for="tool in availableTools" :key="tool" class="tool-switch-item">
|
||||
<span class="tool-name">{{ tool }}</span>
|
||||
<a-switch
|
||||
size="small"
|
||||
:checked="isToolActive(tool)"
|
||||
@change="(checked) => toggleTool(tool, checked)"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
</a-form-item-rest>
|
||||
</a-form-item>
|
||||
|
||||
<!-- 弹窗底部按钮 -->
|
||||
<div class="form-actions" v-if="!state.isEmptyConfig">
|
||||
<a-button type="primary" @click="saveConfig">保存配置</a-button>
|
||||
@ -208,6 +225,7 @@ const router = useRouter();
|
||||
// 状态
|
||||
const agents = ref({});
|
||||
const selectedAgentId = ref(null);
|
||||
const availableTools = ref([]); // 存储所有可用的工具列表
|
||||
const state = reactive({
|
||||
debug_mode: false,
|
||||
isSidebarOpen: JSON.parse(localStorage.getItem('agent-sidebar-open') || 'true'),
|
||||
@ -250,6 +268,47 @@ const closeTokenModal = () => {
|
||||
state.tokenModalVisible = false;
|
||||
};
|
||||
|
||||
// 获取智能体列表
|
||||
const fetchAgents = async () => {
|
||||
try {
|
||||
const response = await fetch('/api/chat/agent');
|
||||
if (response.ok) {
|
||||
const data = await response.json();
|
||||
// 将数组转换为对象
|
||||
agents.value = data.agents.reduce((acc, agent) => {
|
||||
acc[agent.name] = agent;
|
||||
return acc;
|
||||
}, {});
|
||||
// console.log("agents", agents.value);
|
||||
|
||||
// 加载当前选中智能体的配置
|
||||
if (selectedAgentId.value) {
|
||||
loadAgentConfig();
|
||||
}
|
||||
} else {
|
||||
console.error('获取智能体失败');
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('获取智能体错误:', error);
|
||||
}
|
||||
};
|
||||
|
||||
// 获取所有可用工具
|
||||
const fetchTools = async () => {
|
||||
try {
|
||||
const response = await fetch('/api/chat/tools');
|
||||
if (response.ok) {
|
||||
const data = await response.json();
|
||||
availableTools.value = data.tools;
|
||||
console.log("Available tools:", availableTools.value);
|
||||
} else {
|
||||
console.error('获取工具列表失败');
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('获取工具列表错误:', error);
|
||||
}
|
||||
};
|
||||
|
||||
// 根据选中的智能体加载配置
|
||||
const loadAgentConfig = () => {
|
||||
// BUG: 目前消息重置有问题，需要重置消息
|
||||
@ -271,6 +330,10 @@ const loadAgentConfig = () => {
|
||||
agentConfig.value.model = schema.model;
|
||||
}
|
||||
|
||||
if (schema.tools) {
|
||||
agentConfig.value.tools = schema.tools;
|
||||
}
|
||||
|
||||
// 初始化可配置项
|
||||
Object.keys(items).forEach(key => {
|
||||
const item = items[key];
|
||||
@ -303,6 +366,7 @@ const saveConfig = () => {
|
||||
|
||||
// 提示保存成功
|
||||
message.success('配置已保存');
|
||||
console.log("agentConfig.value", agentConfig.value);
|
||||
closeConfigModal();
|
||||
};
|
||||
|
||||
@ -331,31 +395,6 @@ watch(
|
||||
}
|
||||
);
|
||||
|
||||
// 获取智能体列表
|
||||
const fetchAgents = async () => {
|
||||
try {
|
||||
const response = await fetch('/api/chat/agent');
|
||||
if (response.ok) {
|
||||
const data = await response.json();
|
||||
// 将数组转换为对象
|
||||
agents.value = data.agents.reduce((acc, agent) => {
|
||||
acc[agent.name] = agent;
|
||||
return acc;
|
||||
}, {});
|
||||
// console.log("agents", agents.value);
|
||||
|
||||
// 加载当前选中智能体的配置
|
||||
if (selectedAgentId.value) {
|
||||
loadAgentConfig();
|
||||
}
|
||||
} else {
|
||||
console.error('获取智能体失败');
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('获取智能体错误:', error);
|
||||
}
|
||||
};
|
||||
|
||||
// 切换左侧侧边栏
|
||||
const toggleSidebar = () => {
|
||||
state.isSidebarOpen = !state.isSidebarOpen;
|
||||
@ -383,6 +422,8 @@ const selectAgent = (agentId) => {
|
||||
onMounted(async () => {
|
||||
// 获取智能体列表
|
||||
await fetchAgents();
|
||||
// 获取工具列表
|
||||
await fetchTools();
|
||||
|
||||
// 恢复上次选择的智能体
|
||||
const lastSelectedAgent = localStorage.getItem('last-selected-agent');
|
||||
@ -418,6 +459,31 @@ const goToAgentPage = () => {
|
||||
window.open(`/agent/${selectedAgentId.value}`, '_blank');
|
||||
}
|
||||
};
|
||||
|
||||
// 检查工具是否激活
|
||||
const isToolActive = (tool) => {
|
||||
if (!agentConfig.value.tools) {
|
||||
agentConfig.value.tools = [];
|
||||
}
|
||||
return agentConfig.value.tools.includes(tool);
|
||||
};
|
||||
|
||||
// 切换工具状态
|
||||
const toggleTool = (tool, checked) => {
|
||||
if (!agentConfig.value.tools) {
|
||||
agentConfig.value.tools = [];
|
||||
}
|
||||
|
||||
if (checked) {
|
||||
// 添加工具到列表
|
||||
if (!agentConfig.value.tools.includes(tool)) {
|
||||
agentConfig.value.tools.push(tool);
|
||||
}
|
||||
} else {
|
||||
// 从列表中移除工具
|
||||
agentConfig.value.tools = agentConfig.value.tools.filter(item => item !== tool);
|
||||
}
|
||||
};
|
||||
</script>
|
||||
|
||||
<style lang="less" scoped>
|
||||
@ -689,6 +755,22 @@ const goToAgentPage = () => {
|
||||
margin-right: 8px;
|
||||
}
|
||||
}
|
||||
|
||||
.tools-switches {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 12px;
|
||||
|
||||
.tool-switch-item {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: space-between;
|
||||
|
||||
.tool-name {
|
||||
margin-left: 10px;
|
||||
}
|
||||
}
|
||||
}
|
||||
</style>
|
||||
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user