402 lines
13 KiB
Python
Executable File
402 lines
13 KiB
Python
Executable File
import asyncio
|
|
from tools_for_ms import *
|
|
from api_key import *
|
|
from openai import OpenAI
|
|
import json
|
|
from typing import Dict, List, Any, Union, Optional
|
|
from rich.console import Console
|
|
|
|
# 获取工具
|
|
tools = get_tools()
|
|
tools_schema = get_tool_schemas()
|
|
tool_map = {tool_name: tool for tool_name, tool in tools.items()}
|
|
|
|
console = Console()
|
|
console.print(tools_schema)
|
|
|
|
class DualModelAgent:
|
|
"""
|
|
同时支持 qwq 和 gpt4o 两种大模型的代理类
|
|
处理两种不同的返回值格式并提供统一的工具调用接口
|
|
"""
|
|
|
|
def __init__(self, qwq_model_name: str = "qwq-32b", gpt_model_name: str = "gpt-4o"):
|
|
"""
|
|
初始化两个模型的客户端
|
|
|
|
Args:
|
|
qwq_model_name: qwq 模型名称
|
|
gpt_model_name: gpt 模型名称
|
|
"""
|
|
# 初始化 qwq 客户端
|
|
self.qwq_client = OpenAI(
|
|
api_key=OPENAI_API_KEY,
|
|
base_url=OPENAI_API_URL,
|
|
)
|
|
|
|
# 初始化 gpt 客户端 (可以使用不同的 API 密钥和 URL)
|
|
self.gpt_client = OpenAI(
|
|
api_key=OPENAI_API_KEY,
|
|
base_url=OPENAI_API_URL,
|
|
)
|
|
|
|
# 模型名称
|
|
self.qwq_model_name = qwq_model_name
|
|
self.gpt_model_name = gpt_model_name
|
|
|
|
# 定义工具列表
|
|
self.tools = tools_schema
|
|
|
|
def get_qwq_response(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
"""
|
|
获取 qwq 模型的响应（返回字典格式）
|
|
|
|
Args:
|
|
messages: 消息列表
|
|
|
|
Returns:
|
|
字典格式的响应
|
|
"""
|
|
completion = self.qwq_client.chat.completions.create(
|
|
model=self.qwq_model_name,
|
|
messages=messages,
|
|
tools=self.tools,
|
|
temperature=0.6,
|
|
tool_choice='auto'
|
|
)
|
|
# qwq 返回的是对象，需要转换为字典
|
|
return completion.model_dump()
|
|
|
|
def get_gpt_response(self, messages: List[Dict[str, Any]]) -> Any:
|
|
"""
|
|
获取 gpt 模型的响应（返回对象格式）
|
|
|
|
Args:
|
|
messages: 消息列表
|
|
|
|
Returns:
|
|
对象格式的响应
|
|
"""
|
|
completion = self.gpt_client.chat.completions.create(
|
|
model=self.gpt_model_name,
|
|
messages=messages,
|
|
tools=self.tools,
|
|
tool_choice="auto",
|
|
temperature=0.6,
|
|
)
|
|
# gpt 返回的是对象，直接返回
|
|
return completion
|
|
|
|
def extract_tool_calls_from_qwq(self, response: Dict[str, Any]) -> Optional[List[Dict[str, Any]]]:
|
|
"""
|
|
从 qwq 响应中提取工具调用信息
|
|
|
|
Args:
|
|
response: qwq 响应字典
|
|
|
|
Returns:
|
|
工具调用列表，如果没有则返回 None
|
|
"""
|
|
assistant_message = response['choices'][0]['message']
|
|
console.print("assistant_message",assistant_message)
|
|
return assistant_message.get('tool_calls')
|
|
|
|
def extract_tool_calls_from_gpt(self, response: Any) -> Optional[List[Any]]:
|
|
"""
|
|
从 gpt 响应中提取工具调用信息
|
|
|
|
Args:
|
|
response: gpt 响应对象
|
|
|
|
Returns:
|
|
工具调用列表，如果没有则返回 None
|
|
"""
|
|
if hasattr(response.choices[0].message, 'tool_calls'):
|
|
return response.choices[0].message.tool_calls
|
|
return None
|
|
|
|
def extract_content_from_qwq(self, response: Dict[str, Any]) -> str:
|
|
"""
|
|
从 qwq 响应中提取内容
|
|
|
|
Args:
|
|
response: qwq 响应字典
|
|
|
|
Returns:
|
|
内容字符串
|
|
"""
|
|
content = response['choices'][0]['message'].get('content')
|
|
return content if content is not None else ""
|
|
|
|
def extract_content_from_gpt(self, response: Any) -> str:
|
|
"""
|
|
从 gpt 响应中提取内容
|
|
|
|
Args:
|
|
response: gpt 响应对象
|
|
|
|
Returns:
|
|
内容字符串
|
|
"""
|
|
content = response.choices[0].message.content
|
|
return content if content is not None else ""
|
|
|
|
def extract_finish_reason_from_qwq(self, response: Dict[str, Any]) -> str:
|
|
"""
|
|
从 qwq 响应中提取完成原因
|
|
|
|
Args:
|
|
response: qwq 响应字典
|
|
|
|
Returns:
|
|
完成原因
|
|
"""
|
|
return response['choices'][0]['finish_reason']
|
|
|
|
def extract_finish_reason_from_gpt(self, response: Any) -> str:
|
|
"""
|
|
从 gpt 响应中提取完成原因
|
|
|
|
Args:
|
|
response: gpt 响应对象
|
|
|
|
Returns:
|
|
完成原因
|
|
"""
|
|
return response.choices[0].finish_reason
|
|
|
|
async def call_tool(self, tool_name: str, tool_arguments: Dict[str, Any]) -> str:
|
|
"""
|
|
调用工具函数
|
|
|
|
Args:
|
|
tool_name: 工具名称
|
|
tool_arguments: 工具参数
|
|
|
|
Returns:
|
|
工具执行结果
|
|
"""
|
|
if tool_name in tools:
|
|
tool_function = tools[tool_name]
|
|
try:
|
|
# 异步调用工具函数
|
|
tool_result = await tool_function(**tool_arguments)
|
|
return tool_result
|
|
except Exception as e:
|
|
return f"工具调用错误: {str(e)}"
|
|
else:
|
|
return f"未找到工具: {tool_name}"
|
|
|
|
async def chat_with_qwq(self, messages: List[Dict[str, Any]], max_turns: int = 5) -> str:
|
|
"""
|
|
与 qwq 模型对话，支持工具调用
|
|
|
|
Args:
|
|
messages: 初始消息列表
|
|
max_turns: 最大对话轮数
|
|
|
|
Returns:
|
|
最终回答
|
|
"""
|
|
|
|
current_messages = messages.copy()
|
|
turn = 0
|
|
|
|
while turn < max_turns:
|
|
turn += 1
|
|
console.print(f"\n[bold cyan]第 {turn} 轮 QWQ 对话[/bold cyan]")
|
|
console.print("message",current_messages)
|
|
# 获取 qwq 响应
|
|
response = self.get_qwq_response(current_messages)
|
|
assistant_message = response['choices'][0]['message']
|
|
if assistant_message['content'] is None:
|
|
assistant_message['content'] = ""
|
|
console.print("message", assistant_message)
|
|
# 将助手消息添加到上下文
|
|
current_messages.append(assistant_message)
|
|
|
|
# 提取内容和工具调用
|
|
content = self.extract_content_from_qwq(response)
|
|
tool_calls = self.extract_tool_calls_from_qwq(response)
|
|
finish_reason = self.extract_finish_reason_from_qwq(response)
|
|
|
|
console.print(f"[green]助手回复:[/green] {content}")
|
|
|
|
# 如果没有工具调用或已完成，返回内容
|
|
if tool_calls is None or finish_reason != "tool_calls":
|
|
return content
|
|
|
|
# 处理工具调用
|
|
for tool_call in tool_calls:
|
|
tool_call_name = tool_call['function']['name']
|
|
tool_call_arguments = json.loads(tool_call['function']['arguments'])
|
|
|
|
console.print(f"[yellow]调用工具:[/yellow] {tool_call_name}")
|
|
console.print(f"[yellow]工具参数:[/yellow] {tool_call_arguments}")
|
|
|
|
# 执行工具调用
|
|
tool_result = await self.call_tool(tool_call_name, tool_call_arguments)
|
|
console.print(f"[blue]工具结果:[/blue] {tool_result}")
|
|
|
|
# 添加工具结果到上下文
|
|
current_messages.append({
|
|
"role": "tool",
|
|
"name": tool_call_name,
|
|
"content": tool_result
|
|
})
|
|
|
|
return "达到最大对话轮数限制"
|
|
|
|
async def chat_with_gpt(self, messages: List[Dict[str, Any]], max_turns: int = 5) -> str:
|
|
"""
|
|
与 gpt 模型对话，支持工具调用
|
|
|
|
Args:
|
|
messages: 初始消息列表
|
|
max_turns: 最大对话轮数
|
|
|
|
Returns:
|
|
最终回答
|
|
"""
|
|
current_messages = messages.copy()
|
|
turn = 0
|
|
|
|
while turn < max_turns:
|
|
turn += 1
|
|
console.print(f"\n[bold magenta]第 {turn} 轮 GPT 对话[/bold magenta]")
|
|
|
|
# 获取 gpt 响应
|
|
response = self.get_gpt_response(current_messages)
|
|
assistant_message = response.choices[0].message
|
|
|
|
# 将助手消息添加到上下文
|
|
current_messages.append(assistant_message.model_dump())
|
|
|
|
# 提取内容和工具调用
|
|
content = self.extract_content_from_gpt(response)
|
|
tool_calls = self.extract_tool_calls_from_gpt(response)
|
|
finish_reason = self.extract_finish_reason_from_gpt(response)
|
|
|
|
console.print(f"[green]助手回复:[/green] {content}")
|
|
|
|
# 如果没有工具调用或已完成，返回内容
|
|
if tool_calls is None or finish_reason != "tool_calls":
|
|
return content
|
|
|
|
# 处理工具调用
|
|
for tool_call in tool_calls:
|
|
tool_call_name = tool_call.function.name
|
|
tool_call_arguments = json.loads(tool_call.function.arguments)
|
|
|
|
console.print(f"[yellow]调用工具:[/yellow] {tool_call_name}")
|
|
console.print(f"[yellow]工具参数:[/yellow] {tool_call_arguments}")
|
|
|
|
# 执行工具调用
|
|
tool_result = await self.call_tool(tool_call_name, tool_call_arguments)
|
|
console.print(f"[blue]工具结果:[/blue] {tool_result}")
|
|
|
|
# 添加工具结果到上下文
|
|
current_messages.append({
|
|
"role": "tool",
|
|
"tool_call_id": tool_call.id,
|
|
"name": tool_call_name,
|
|
"content": tool_result
|
|
})
|
|
|
|
return "达到最大对话轮数限制"
|
|
|
|
async def chat_with_both_models(self, user_input: str, system_prompt: str = "你是一个有用的助手。") -> Dict[str, str]:
|
|
"""
|
|
同时与两个模型对话，比较它们的回答
|
|
|
|
Args:
|
|
user_input: 用户输入
|
|
system_prompt: 系统提示
|
|
|
|
Returns:
|
|
包含两个模型回答的字典
|
|
"""
|
|
messages = [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_input}
|
|
]
|
|
|
|
# 并行调用两个模型
|
|
qwq_task = asyncio.create_task(self.chat_with_qwq(messages))
|
|
gpt_task = asyncio.create_task(self.chat_with_gpt(messages))
|
|
|
|
# 等待两个任务完成
|
|
qwq_response, gpt_response = await asyncio.gather(qwq_task, gpt_task)
|
|
|
|
return {
|
|
"qwq": qwq_response,
|
|
"gpt": gpt_response
|
|
}
|
|
|
|
|
|
async def main():
|
|
"""主函数"""
|
|
# 创建双模型代理
|
|
agent = DualModelAgent()
|
|
|
|
while True:
|
|
# 获取用户输入
|
|
user_input = input("\n请输入问题（输入 'exit' 退出）: ")
|
|
#user_input='现在的时间'
|
|
if user_input.lower() == 'exit':
|
|
break
|
|
|
|
# 选择模型
|
|
model_choice = input("选择模型 (1: QWQ, 2: GPT, 3: 两者): ")
|
|
|
|
|
|
try:
|
|
if model_choice == '1':
|
|
# 使用 QWQ 模型
|
|
messages = [
|
|
{"role": "system", "content": "你是一个有用的助手。"},
|
|
{"role": "user", "content": user_input}
|
|
]
|
|
response = await agent.chat_with_qwq(messages)
|
|
console.print(f"\n[bold cyan]QWQ 最终回答:[/bold cyan] {response}")
|
|
|
|
elif model_choice == '2':
|
|
# 使用 GPT 模型
|
|
messages = [
|
|
{"role": "system", "content": "你是一个有用的助手。"},
|
|
{"role": "user", "content": user_input}
|
|
]
|
|
response = await agent.chat_with_gpt(messages)
|
|
console.print(f"\n[bold magenta]GPT 最终回答:[/bold magenta] {response}")
|
|
|
|
elif model_choice == '3':
|
|
# 同时使用两个模型
|
|
responses = await agent.chat_with_both_models(user_input)
|
|
console.print(f"\n[bold cyan]QWQ 最终回答:[/bold cyan] {responses['qwq']}")
|
|
console.print(f"\n[bold magenta]GPT 最终回答:[/bold magenta] {responses['gpt']}")
|
|
|
|
else:
|
|
console.print("[bold red]无效的选择，请输入 1、2 或 3[/bold red]")
|
|
|
|
except Exception as e:
|
|
console.print(f"[bold red]错误:[/bold red] {str(e)}")
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
asyncio.run(main())
|
|
#from tools_for_ms.services_tools.mp_service_tools import search_material_property_from_material_project
|
|
#asyncio.run(search_material_property_from_material_project('Fe2O3'))
|
|
|
|
# 测试工具函数
|
|
|
|
# tool_name = 'get_crystal_structures_from_materials_project'
|
|
# result = asyncio.run(test_tool(tool_name))
|
|
# print(result)
|
|
|
|
#pass
|
|
|
|
# 知识检索API的接口 数据库
|