kris/ForcePilot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Merge pull request #135 from xerrors/132-feat-retrieve-as-tool

Browse Source 
添加自定义工具功能，给未来的 MCP 打一下基础
This commit is contained in:
Wenjie Zhang 2025-04-05 17:30:34 +08:00 committed by GitHub
commit bde802f64a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 208 additions and 88 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
server/routers/base_router.py
View File

@ -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!"}

6
server/routers/chat_router.py
View File

@ -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())}

9
src/agents/chatbot/configuration.py
View File

@ -34,3 +34,12 @@ class ChatbotConfiguration(Configuration):
"description": "智能体的驱动模型"
},
)
tools: list[str] = field(
default_factory=list,
metadata={
"name": "工具",
"configurable": False,
"description": "工具列表"
},
)

36
src/agents/chatbot/graph.py
View File

@ -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",

20
src/agents/react/workflows.py
View File

@ -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

63
src/agents/tools_factory.py
View File

@ -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,
}

13
src/core/graphbase.py
View File

@ -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()

12
src/core/knowledgebase.py
View File

@ -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 #
################################

3
web/src/components/AgentChatComponent.vue
View File

@ -360,6 +360,7 @@ const sendMessageWithText = async (text) => {
body: JSON.stringify(requestData)
});
// console.log("requestData", requestData);
if (!response.ok) {
throw new Error('请求失败');
}
@ -554,7 +555,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);

132
web/src/views/AgentView.vue
View File

@ -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>