import asyncio import json import sys import os from openai import OpenAI from rich.console import Console from rich.panel import Panel from rich.text import Text # 添加项目根目录到sys.path sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from sci_mcp import * # 创建Rich控制台对象 console = Console() # 创建一个列表来存储工具调用结果 tool_results = [] # 定义分隔符样式 def print_separator(title=""): console.print(Panel(f"[bold magenta]{title}[/]", border_style="cyan", expand=False), justify="center") api_key="gpustack_41f8963fb74e9f39_931a02e3ab35b41b1eada85da9c40705" base_url="http://gpustack.ddwtop.team/v1-openai" client = OpenAI( api_key=api_key, base_url=base_url, ) from sci_mcp import * tools_schemas = get_domain_tool_schemas(["material",'general']) tool_map = get_domain_tools(["material",'general']) # 打印消息的函数 def print_message(message): # 处理不同类型的消息对象 if hasattr(message, 'role'): # ChatCompletionMessage 对象 role = message.role content = message.content if hasattr(message, 'content') else "" # 如果是工具消息，获取工具名称 tool_name = None if role == "tool" and hasattr(message, 'name'): tool_name = message.name else: # 字典类型 role = message.get("role", "unknown") content = message.get("content", "") # 如果是工具消息，获取工具名称 tool_name = message.get("name") if role == "tool" else None # 根据角色选择不同的颜色 role_colors = { "system": "bright_blue", "user": "green", "assistant": "yellow", "tool": "bright_red" } color = role_colors.get(role, "white") # 创建富文本面板 text = Text() # 如果是工具消息，添加工具名称 if role == "tool" and tool_name: text.append(f"{role} ({tool_name}): ", style=f"bold {color}") else: text.append(f"{role}: ", style=f"bold {color}") text.append(str(content)) console.print(Panel(text, border_style=color)) messages = [ {"role": "user", "content": """data_Ti4V _symmetry_space_group_name_H-M Fmmm _cell_length_a 3.18353600 _cell_length_b 4.52677200 _cell_length_c 22.74397000 _cell_angle_alpha 90.00000000 _cell_angle_beta 90.00000000 _cell_angle_gamma 90.00000000 _symmetry_Int_Tables_number 69 _chemical_formula_structural Ti4V _chemical_formula_sum 'Ti16 V4' _cell_volume 327.76657340 _cell_formula_units_Z 4 loop_ _symmetry_equiv_pos_site_id _symmetry_equiv_pos_as_xyz 1 'x, y, z' 2 '-x, -y, -z' 3 '-x, -y, z' 4 'x, y, -z' 5 'x, -y, -z' 6 '-x, y, z' 7 '-x, y, -z' 8 'x, -y, z' 9 'x+1/2, y, z+1/2' 10 '-x+1/2, -y, -z+1/2' 11 '-x+1/2, -y, z+1/2' 12 'x+1/2, y, -z+1/2' 13 'x+1/2, -y, -z+1/2' 14 '-x+1/2, y, z+1/2' 15 '-x+1/2, y, -z+1/2' 16 'x+1/2, -y, z+1/2' 17 'x+1/2, y+1/2, z' 18 '-x+1/2, -y+1/2, -z' 19 '-x+1/2, -y+1/2, z' 20 'x+1/2, y+1/2, -z' 21 'x+1/2, -y+1/2, -z' 22 '-x+1/2, y+1/2, z' 23 '-x+1/2, y+1/2, -z' 24 'x+1/2, -y+1/2, z' 25 'x, y+1/2, z+1/2' 26 '-x, -y+1/2, -z+1/2' 27 '-x, -y+1/2, z+1/2' 28 'x, y+1/2, -z+1/2' 29 'x, -y+1/2, -z+1/2' 30 '-x, y+1/2, z+1/2' 31 '-x, y+1/2, -z+1/2' 32 'x, -y+1/2, z+1/2' loop_ _atom_site_type_symbol _atom_site_label _atom_site_symmetry_multiplicity _atom_site_fract_x _atom_site_fract_y _atom_site_fract_z _atom_site_occupancy Ti Ti0 8 0.00000000 0.00000000 0.19939400 1.0 Ti Ti1 8 0.00000000 0.00000000 0.40878900 1.0 V V2 4 0.00000000 0.00000000 0.00000000 1.0 ，根据上文提供的CIF文件，请你分析该晶体材料的化学式、晶体系统、空间群、密度、是否为金属、磁性和稳定性，并用JSON格式回答。"""} ] #how to synthesize CsPbBr3 at room temperature # # 打印初始消息 print_separator("初始消息") for message in messages: print_message(message) finish_reason = None async def execute_tool(tool_name,tool_arguments): # 设置LLM调用上下文标记 try: tool_func = tool_map[tool_name] # 获取工具函数 arguments = {} if tool_arguments: # 检查arguments是字符串还是字典 if isinstance(tool_arguments, dict): # 如果已经是字典，直接使用 arguments = tool_arguments elif isinstance(tool_arguments, str): # 如果是字符串，尝试解析为JSON try: # 尝试直接解析为JSON对象 arguments = json.loads(tool_arguments) except json.JSONDecodeError: # 如果解析失败，可能是因为字符串中包含转义字符 # 尝试修复常见的JSON字符串问题 fixed_str = tool_arguments.replace('\\"', '"').replace('\\\\', '\\') try: arguments = json.loads(fixed_str) except json.JSONDecodeError: # 如果仍然失败，尝试将字符串作为原始字符串处理 arguments = {"raw_string": tool_arguments} # 调用工具函数 if asyncio.iscoroutinefunction(tool_func): # 如果是异步函数，使用await调用 result = await tool_func(**arguments) else: # 如果是同步函数，直接调用 result = tool_func(**arguments) # if func_name=='generate_material': # print("xxxxx",result) return result finally: # 清除LLM调用上下文标记 pass # 定义一个函数来估算消息的token数量（粗略估计） def estimate_tokens(messages): # 简单估计：每个英文单词约1.3个token，每个中文字符约1个token total = 0 for msg in messages: content = msg.get("content", "") if isinstance(msg, dict) else (msg.content if hasattr(msg, "content") else "") if content: # 粗略估计内容的token数 total += len(content) * 1.3 # 估计工具调用的token数 tool_calls = msg.get("tool_calls", []) if isinstance(msg, dict) else (msg.tool_calls if hasattr(msg, "tool_calls") else []) for tool_call in tool_calls: if isinstance(tool_call, dict): args = tool_call.get("function", {}).get("arguments", "") total += len(args) * 1.3 else: args = tool_call.function.arguments if hasattr(tool_call, "function") else "" total += len(args) * 1.3 return int(total) # 管理消息历史，保持在token限制以内 def manage_message_history(messages, max_tokens=7000): # 保留第一条消息（通常是系统消息或初始用户消息） if len(messages) <= 1: return messages # 估算当前消息的token数 current_tokens = estimate_tokens(messages) # 如果当前token数已经接近限制，开始裁剪历史消息 if current_tokens > max_tokens: # 保留第一条消息和最近的消息 preserved_messages = [messages[0]] # 从最新的消息开始添加，直到接近但不超过token限制 temp_messages = [] for msg in reversed(messages[1:]): temp_messages.insert(0, msg) if estimate_tokens([messages[0]] + temp_messages) > max_tokens: # 如果添加这条消息会超过限制，则停止添加 temp_messages.pop(0) break # 如果裁剪后的消息太少，至少保留最近的几条消息 if len(temp_messages) < 4 and len(messages) > 4: temp_messages = messages[-4:] return preserved_messages + temp_messages return messages while finish_reason is None or finish_reason == "tool_calls": # 在发送请求前管理消息历史 managed_messages = manage_message_history(messages) if len(managed_messages) < len(messages): print_separator(f"消息历史已裁剪，从{len(messages)}条减少到{len(managed_messages)}条") messages = managed_messages completion = client.chat.completions.create( model="MARS-T1", messages=messages, temperature=0.3, tools=tools_schemas, timeout=120, ) choice = completion.choices[0] finish_reason = choice.finish_reason if finish_reason == "tool_calls": # <-- 判断当前返回内容是否包含 tool_calls # 打印assistant消息 print_separator("Assistant消息") print_message(choice.message) # 将ChatCompletionMessage对象转换为字典 assistant_message = { "role": "assistant", "content": choice.message.content if hasattr(choice.message, 'content') else None } # 如果有工具调用，添加到字典中 if hasattr(choice.message, 'tool_calls') and choice.message.tool_calls: # 将tool_calls对象转换为字典列表 tool_calls_list = [] for tool_call in choice.message.tool_calls: tool_call_dict = { "id": tool_call.id, "type": "function", "function": { "name": tool_call.function.name, "arguments": tool_call.function.arguments } } tool_calls_list.append(tool_call_dict) assistant_message["tool_calls"] = tool_calls_list # 添加消息到上下文 messages.append(assistant_message) # <-- 我们将模型返回给我们的 assistant 消息也添加到上下文中，以便于下次请求时模型能理解我们的诉求 # 打印工具调用信息 print_separator("工具调用") for tool_call in choice.message.tool_calls: console.print(f"[bold cyan]工具名称:[/] [yellow]{tool_call.function.name}[/]") console.print(f"[bold cyan]工具ID:[/] [yellow]{tool_call.id}[/]") console.print(f"[bold cyan]参数:[/] [yellow]{tool_call.function.arguments}[/]") console.print("") tool_call_name = tool_call.function.name tool_call_arguments = json.loads(tool_call.function.arguments) # <-- arguments 是序列化后的 JSON Object，我们需要使用 json.loads 反序列化一下 try: tool_result = asyncio.run(execute_tool(tool_name=tool_call_name,tool_arguments=tool_call_arguments)) # <-- 通过 tool_map 快速找到需要执行哪个函数 except Exception as e: tool_result=f'工具调用失败{e}' # 将工具调用结果保存到单独的列表中，使用指定格式包裹 formatted_result = f"[{tool_call_name} content begin]{tool_result}[{tool_call_name} content end]" tool_results.append({ "tool_name": tool_call_name, "tool_id": tool_call.id, "formatted_result": formatted_result, "raw_result": tool_result }) # 打印保存的工具结果信息 console.print(f"[bold green]已保存工具结果:[/] [cyan]{tool_call_name}[/]") # 构造工具响应消息 tool_message = { "role": "tool", "tool_call_id": tool_call.id, "name": tool_call_name, "content": tool_result, # <-- 我们约定使用字符串格式向模型提交工具调用结果 } # 打印工具响应 print_separator(f"工具响应: {tool_call_name}") print_message(tool_message) # 添加消息到上下文 #messages.append(tool_message) # 打印最终响应 if choice.message.content: print_separator("最终响应") console.print(Panel(choice.message.content, border_style="green")) # 打印收集的所有工具调用结果 if tool_results: print_separator("所有工具调用结果") console.print(f"[bold cyan]共收集了 {len(tool_results)} 个工具调用结果[/]") # 将所有格式化的结果写入文件 with open("tool_results.txt", "w", encoding="utf-8") as f: for result in tool_results: f.write(f"{result['formatted_result']}\n\n") console.print(f"[bold green]所有工具调用结果已保存到 tool_results.txt 文件中[/]")