343 lines
13 KiB
Python
343 lines
13 KiB
Python
# Copyright 2023 The Qwen team, Alibaba Group. All rights reserved.
|
||
#
|
||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
# you may not use this file except in compliance with the License.
|
||
# You may obtain a copy of the License at
|
||
#
|
||
# http://www.apache.org/licenses/LICENSE-2.0
|
||
#
|
||
# Unless required by applicable law or agreed to in writing, software
|
||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
# See the License for the specific language governing permissions and
|
||
# limitations under the License.
|
||
|
||
"""A sqlite database assistant implemented by assistant"""
|
||
|
||
import argparse
|
||
import asyncio
|
||
import copy
|
||
import json
|
||
import os
|
||
from typing import Dict, List, Optional, Union
|
||
|
||
from qwen_agent.agents import Assistant
|
||
from qwen_agent.gui import WebUI
|
||
from qwen_agent.llm.oai import TextChatAtOAI
|
||
from qwen_agent.llm.schema import ASSISTANT, FUNCTION, Message
|
||
from qwen_agent.utils.output_beautify import typewriter_print
|
||
|
||
ROOT_RESOURCE = os.path.join(os.path.dirname(__file__), "resource")
|
||
|
||
|
||
class GPT4OChat(TextChatAtOAI):
|
||
"""自定义 GPT-4o 聊天类，修复 tool_call_id 问题"""
|
||
|
||
def convert_messages_to_dicts(self, messages: List[Message]) -> List[dict]:
|
||
# 使用父类方法进行基础转换
|
||
messages = super().convert_messages_to_dicts(messages)
|
||
|
||
# 应用修复后的消息转换
|
||
messages = self._fixed_conv_qwen_agent_messages_to_oai(messages)
|
||
|
||
return messages
|
||
|
||
@staticmethod
|
||
def _fixed_conv_qwen_agent_messages_to_oai(messages: List[Union[Message, Dict]]):
|
||
"""修复后的消息转换方法，确保 tool 消息包含 tool_call_id 字段"""
|
||
new_messages = []
|
||
i = 0
|
||
|
||
while i < len(messages):
|
||
msg = messages[i]
|
||
|
||
if msg['role'] == ASSISTANT:
|
||
# 处理 assistant 消息
|
||
assistant_msg = {'role': 'assistant'}
|
||
|
||
# 设置 content
|
||
content = msg.get('content', '')
|
||
if isinstance(content, (list, dict)):
|
||
assistant_msg['content'] = json.dumps(content, ensure_ascii=False)
|
||
elif content is None:
|
||
assistant_msg['content'] = ''
|
||
else:
|
||
assistant_msg['content'] = content
|
||
|
||
# 设置 reasoning_content
|
||
if msg.get('reasoning_content'):
|
||
assistant_msg['reasoning_content'] = msg['reasoning_content']
|
||
|
||
# 检查是否需要构造 tool_calls
|
||
has_tool_call = False
|
||
tool_calls = []
|
||
|
||
# 情况1：当前消息有 function_call
|
||
if msg.get('function_call'):
|
||
has_tool_call = True
|
||
tool_calls.append({
|
||
'id': msg.get('extra', {}).get('function_id', '1'),
|
||
'type': 'function',
|
||
'function': {
|
||
'name': msg['function_call']['name'],
|
||
'arguments': msg['function_call']['arguments']
|
||
}
|
||
})
|
||
|
||
# 注意：不再为孤立的 tool 消息构造虚假的 tool_call
|
||
|
||
if has_tool_call:
|
||
assistant_msg['tool_calls'] = tool_calls
|
||
new_messages.append(assistant_msg)
|
||
|
||
# 检查后续是否有对应的 tool 消息
|
||
if i + 1 < len(messages) and messages[i + 1]['role'] == 'tool':
|
||
tool_msg = copy.deepcopy(messages[i + 1])
|
||
# 确保 tool_call_id 匹配
|
||
tool_msg['tool_call_id'] = tool_calls[0]['id']
|
||
# 移除多余字段
|
||
for field in ['id', 'extra', 'function_call']:
|
||
if field in tool_msg:
|
||
del tool_msg[field]
|
||
# 确保 content 有效且为字符串
|
||
content = tool_msg.get('content', '')
|
||
if isinstance(content, (list, dict)):
|
||
tool_msg['content'] = json.dumps(content, ensure_ascii=False)
|
||
elif content is None:
|
||
tool_msg['content'] = ''
|
||
new_messages.append(tool_msg)
|
||
i += 2
|
||
else:
|
||
i += 1
|
||
else:
|
||
new_messages.append(assistant_msg)
|
||
i += 1
|
||
|
||
elif msg['role'] == 'tool':
|
||
# 孤立的 tool 消息，转换为 assistant + user 消息序列
|
||
# 首先添加一个包含工具结果的 assistant 消息
|
||
assistant_result = {'role': 'assistant'}
|
||
content = msg.get('content', '')
|
||
if isinstance(content, (list, dict)):
|
||
content = json.dumps(content, ensure_ascii=False)
|
||
assistant_result['content'] = f"工具查询结果: {content}"
|
||
new_messages.append(assistant_result)
|
||
|
||
# 然后添加一个 user 消息来继续对话
|
||
new_messages.append({'role': 'user', 'content': '请继续分析以上结果'})
|
||
i += 1
|
||
|
||
else:
|
||
# 处理其他角色消息
|
||
new_msg = copy.deepcopy(msg)
|
||
|
||
# 确保 content 有效且为字符串
|
||
content = new_msg.get('content', '')
|
||
if isinstance(content, (list, dict)):
|
||
new_msg['content'] = json.dumps(content, ensure_ascii=False)
|
||
elif content is None:
|
||
new_msg['content'] = ''
|
||
|
||
new_messages.append(new_msg)
|
||
i += 1
|
||
|
||
return new_messages
|
||
|
||
|
||
def read_mcp_settings():
|
||
with open("./mcp/mcp_settings.json", "r") as f:
|
||
mcp_settings_json = json.load(f)
|
||
return mcp_settings_json
|
||
|
||
|
||
def read_system_prompt():
|
||
with open("./agent_prompt.txt", "r", encoding="utf-8") as f:
|
||
return f.read().strip()
|
||
|
||
|
||
def init_agent_service():
|
||
llm_cfg = {
|
||
"llama-33": {
|
||
"model": "gbase-llama-33",
|
||
"model_server": "http://llmapi:9009/v1",
|
||
"api_key": "any",
|
||
},
|
||
"gpt-oss-120b": {
|
||
"model": "openai/gpt-oss-120b",
|
||
"model_server": "https://openrouter.ai/api/v1", # base_url, also known as api_base
|
||
"api_key": "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",
|
||
|
||
"generate_cfg": {
|
||
"use_raw_api": True, # GPT-OSS true ,Qwen false
|
||
"fncall_prompt_type": "nous", # 使用 nous 风格的函数调用提示
|
||
}
|
||
},
|
||
|
||
"claude-3.7": {
|
||
"model": "claude-3-7-sonnet-20250219",
|
||
"model_server": "https://one.felo.me/v1",
|
||
"api_key": "sk-9gtHriq7C3jAvepq5dA0092a5cC24a54Aa83FbC99cB88b21-2",
|
||
"generate_cfg": {
|
||
"use_raw_api": True, # GPT-OSS true ,Qwen false
|
||
},
|
||
},
|
||
"gpt-4o": {
|
||
"model": "gpt-4o",
|
||
"model_server": "https://one-dev.felo.me/v1",
|
||
"api_key": "sk-hsKClH0Z695EkK5fDdB2Ec2fE13f4fC1B627BdBb8e554b5b-4",
|
||
"generate_cfg": {
|
||
"use_raw_api": True, # 启用 raw_api 但使用自定义类修复 tool_call_id 问题
|
||
"fncall_prompt_type": "nous", # 使用 nous 风格的函数调用提示
|
||
},
|
||
},
|
||
"Gpt-4o-back": {
|
||
"model_type": "oai", # 使用 oai 类型以便使用自定义类
|
||
"model": "gpt-4o",
|
||
"model_server": "https://one-dev.felo.me/v1",
|
||
"api_key": "sk-hsKClH0Z695EkK5fDdB2Ec2fE13f4fC1B627BdBb8e554b5b-4",
|
||
"generate_cfg": {
|
||
"use_raw_api": True, # 启用 raw_api 但使用自定义类修复 tool_call_id 问题
|
||
"fncall_prompt_type": "nous", # 使用 nous 风格的函数调用提示
|
||
},
|
||
# 使用自定义的 GPT4OChat 类
|
||
"llm_class": GPT4OChat,
|
||
},
|
||
|
||
"glm-45": {
|
||
"model_server": "https://open.bigmodel.cn/api/paas/v4",
|
||
"api_key": "0c9cbaca9d2bbf864990f1e1decdf340.dXRMsZCHTUbPQ0rm",
|
||
"model": "glm-4.5",
|
||
"generate_cfg": {
|
||
"use_raw_api": True, # GPT-OSS true ,Qwen false
|
||
},
|
||
},
|
||
"qwen3-next": {
|
||
"model": "qwen/qwen3-next-80b-a3b-instruct",
|
||
"model_server": "https://openrouter.ai/api/v1", # base_url, also known as api_base
|
||
"api_key": "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",
|
||
|
||
},
|
||
"deepresearch": {
|
||
"model": "alibaba/tongyi-deepresearch-30b-a3b",
|
||
"model_server": "https://openrouter.ai/api/v1", # base_url, also known as api_base
|
||
"api_key": "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",
|
||
},
|
||
|
||
"qwen3-coder":{
|
||
"model": "Qwen/Qwen3-Coder-30B-A3B-Instruct",
|
||
"model_server": "https://api-inference.modelscope.cn/v1", # base_url, also known as api_base
|
||
"api_key": "ms-92027446-2787-4fd6-af01-f002459ec556",
|
||
},
|
||
"openrouter-gpt4o":{
|
||
"model": "openai/gpt-4o",
|
||
"model_server": "https://openrouter.ai/api/v1", # base_url, also known as api_base
|
||
"api_key": "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",
|
||
"generate_cfg": {
|
||
"use_raw_api": True, # GPT-OSS true ,Qwen false
|
||
"fncall_prompt_type": "nous", # 使用 nous 风格的函数调用提示
|
||
},
|
||
}
|
||
}
|
||
system = read_system_prompt()
|
||
|
||
# 暂时禁用 MCP 工具以测试 GPT-4o
|
||
tools = read_mcp_settings()
|
||
# 使用自定义的 GPT-4o 配置
|
||
llm_instance = llm_cfg["qwen3-next"]
|
||
if "llm_class" in llm_instance:
|
||
llm_instance = llm_instance.get("llm_class", TextChatAtOAI)(llm_instance)
|
||
|
||
bot = Assistant(
|
||
llm=llm_instance, # 使用自定义的 GPT-4o 实例
|
||
name="数据库助手",
|
||
description="数据库查询",
|
||
system_message=system,
|
||
function_list=tools,
|
||
)
|
||
|
||
return bot
|
||
|
||
|
||
def test(query="数据库里有几张表"):
|
||
# Define the agent
|
||
bot = init_agent_service()
|
||
|
||
# Chat
|
||
messages = []
|
||
|
||
messages.append({"role": "user", "content": query})
|
||
|
||
responses = []
|
||
for response in bot.run(messages):
|
||
responses.append(response)
|
||
|
||
# 只输出最终结果，不显示中间过程
|
||
if responses:
|
||
final_response = responses[-1][-1] # 取最后一个响应作为最终结果
|
||
print("Answer:", final_response["content"])
|
||
|
||
|
||
def app_tui():
|
||
# Define the agent
|
||
bot = init_agent_service()
|
||
|
||
# Chat
|
||
messages = []
|
||
while True:
|
||
# Query example: 数据库里有几张表
|
||
query = input("user question: ")
|
||
# File example: resource/poem.pdf
|
||
file = input("file url (press enter if no file): ").strip()
|
||
if not query:
|
||
print("user question cannot be empty！")
|
||
continue
|
||
if not file:
|
||
messages.append({"role": "user", "content": query})
|
||
else:
|
||
messages.append(
|
||
{"role": "user", "content": [{"text": query}, {"file": file}]}
|
||
)
|
||
|
||
response = []
|
||
for response in bot.run(messages):
|
||
print("bot response:", response)
|
||
messages.extend(response)
|
||
|
||
|
||
def app_gui():
|
||
# Define the agent
|
||
bot = init_agent_service()
|
||
chatbot_config = {
|
||
"prompt.suggestions": [
|
||
"数据库里有几张表",
|
||
"创建一个学生表包括学生的姓名、年龄",
|
||
"增加一个学生名字叫韩梅梅，今年6岁",
|
||
]
|
||
}
|
||
WebUI(
|
||
bot,
|
||
chatbot_config=chatbot_config,
|
||
).run()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(description="数据库助手")
|
||
parser.add_argument(
|
||
"--query", type=str, default="数据库里有几张表", help="用户问题"
|
||
)
|
||
parser.add_argument(
|
||
"--mode",
|
||
type=str,
|
||
choices=["test", "tui", "gui"],
|
||
default="test",
|
||
help="运行模式",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
if args.mode == "test":
|
||
test(args.query)
|
||
elif args.mode == "tui":
|
||
app_tui()
|
||
elif args.mode == "gui":
|
||
app_gui()
|