290 lines
12 KiB
Python
290 lines
12 KiB
Python
"""
|
||
测试RXN工具函数模块
|
||
|
||
此模块包含用于测试rxn_tools模块中化学反应预测和分析工具函数的测试用例。
|
||
"""
|
||
import asyncio
|
||
import sys
|
||
import os
|
||
from pathlib import Path
|
||
from rich.console import Console
|
||
|
||
# 添加项目根目录到Python路径
|
||
sys.path.append('/home/ubuntu/sas0/lzy/multi_mcp_server')
|
||
|
||
from sci_mcp.chemistry_mcp.rxn_tools.rxn_tools import (
|
||
predict_reaction_outcome_rxn,
|
||
predict_reaction_topn_rxn,
|
||
predict_reaction_properties_rxn,
|
||
extract_reaction_actions_rxn
|
||
)
|
||
|
||
# 创建控制台对象用于格式化输出
|
||
console = Console()
|
||
|
||
|
||
async def test_predict_reaction_outcome():
|
||
"""测试反应结果预测功能"""
|
||
console.print("[bold cyan]测试反应结果预测功能[/bold cyan]")
|
||
|
||
# 使用固定参数：溴和蒽的反应
|
||
fixed_reactants = "BrBr.c1ccc2cc3ccccc3cc2c1"
|
||
console.print(f"固定反应物: {fixed_reactants}")
|
||
|
||
try:
|
||
result = await predict_reaction_outcome(fixed_reactants)
|
||
console.print("[green]结果:[/green]")
|
||
console.print(result)
|
||
except Exception as e:
|
||
console.print(f"[bold red]错误:[/bold red] {str(e)}")
|
||
|
||
|
||
async def test_predict_reaction_topn():
|
||
"""测试多产物预测功能"""
|
||
console.print("\n[bold cyan]测试多产物预测功能[/bold cyan]")
|
||
|
||
# 测试1：单个反应（字符串格式）
|
||
fixed_reactants = "C=CC=O.CN" # 丙烯醛和甲胺
|
||
fixed_topn = 2
|
||
console.print(f"测试1 - 单个反应（字符串格式）")
|
||
console.print(f"固定反应物: {fixed_reactants}")
|
||
console.print(f"固定预测产物数量: {fixed_topn}")
|
||
|
||
try:
|
||
result = await predict_reaction_topn(fixed_reactants, fixed_topn)
|
||
console.print("[green]结果:[/green]")
|
||
console.print(result)
|
||
except Exception as e:
|
||
console.print(f"[bold red]错误:[/bold red] {str(e)}")
|
||
|
||
# 测试2：单个反应（列表格式）
|
||
fixed_reactants_list = ["BrBr", "c1ccc2cc3ccccc3cc2c1"] # 溴和蒽
|
||
console.print(f"\n测试2 - 单个反应（列表格式）")
|
||
console.print(f"固定反应物: {fixed_reactants_list}")
|
||
console.print(f"固定预测产物数量: {fixed_topn}")
|
||
|
||
try:
|
||
result = await predict_reaction_topn(fixed_reactants_list, fixed_topn)
|
||
console.print("[green]结果:[/green]")
|
||
console.print(result)
|
||
except Exception as e:
|
||
console.print(f"[bold red]错误:[/bold red] {str(e)}")
|
||
|
||
# 测试3：多个反应（列表的列表格式）
|
||
fixed_reactants_batch = [
|
||
["BrBr", "c1ccc2cc3ccccc3cc2c1"], # 溴和蒽
|
||
["BrBr", "c1ccc2cc3ccccc3cc2c1CCO"] # 溴和修饰的蒽
|
||
]
|
||
console.print(f"\n测试3 - 多个反应（列表的列表格式）")
|
||
console.print(f"固定反应物批量: {fixed_reactants_batch}")
|
||
console.print(f"固定预测产物数量: {fixed_topn}")
|
||
|
||
try:
|
||
result = await predict_reaction_topn(fixed_reactants_batch, fixed_topn)
|
||
console.print("[green]结果:[/green]")
|
||
console.print(result)
|
||
except Exception as e:
|
||
console.print(f"[bold red]错误:[/bold red] {str(e)}")
|
||
console.print(f"[bold red]错误:[/bold red] {str(e)}")
|
||
|
||
|
||
async def test_predict_retrosynthesis():
|
||
"""测试逆合成分析功能"""
|
||
console.print("\n[bold cyan]测试逆合成分析功能[/bold cyan]")
|
||
|
||
# 使用固定参数：阿司匹林的逆合成分析
|
||
fixed_target_molecule = "CC(=O)OC1=CC=CC=C1C(=O)O" # 阿司匹林
|
||
fixed_max_steps = 1
|
||
console.print(f"固定目标分子: {fixed_target_molecule}")
|
||
console.print(f"固定最大步骤数: {fixed_max_steps}")
|
||
|
||
try:
|
||
result = await predict_retrosynthesis(fixed_target_molecule, fixed_max_steps)
|
||
console.print("[green]结果:[/green]")
|
||
console.print(result)
|
||
except Exception as e:
|
||
console.print(f"[bold red]错误:[/bold red] {str(e)}")
|
||
|
||
|
||
async def test_predict_biocatalytic_retrosynthesis():
|
||
"""测试生物催化逆合成分析功能"""
|
||
console.print("\n[bold cyan]测试生物催化逆合成分析功能[/bold cyan]")
|
||
|
||
# 使用固定参数：一个可能适合酶催化的分子
|
||
fixed_target_molecule = "OC1C(O)C=C(Br)C=C1"
|
||
console.print(f"固定目标分子: {fixed_target_molecule}")
|
||
|
||
try:
|
||
result = await predict_biocatalytic_retrosynthesis(fixed_target_molecule)
|
||
console.print("[green]结果:[/green]")
|
||
console.print(result)
|
||
except Exception as e:
|
||
console.print(f"[bold red]错误:[/bold red] {str(e)}")
|
||
|
||
|
||
async def test_predict_reaction_properties():
|
||
"""测试反应属性预测功能"""
|
||
console.print("\n[bold cyan]测试反应属性预测功能[/bold cyan]")
|
||
|
||
# 使用固定参数：原子映射
|
||
fixed_reaction = "CC(C)S.CN(C)C=O.Fc1cccnc1F.O=C([O-])[O-].[K+].[K+]>>CC(C)Sc1ncccc1F"
|
||
fixed_property_type = "atom-mapping"
|
||
console.print(f"固定反应: {fixed_reaction}")
|
||
console.print(f"固定属性类型: {fixed_property_type}")
|
||
|
||
try:
|
||
result = await predict_reaction_properties(fixed_reaction, fixed_property_type)
|
||
console.print("[green]结果:[/green]")
|
||
console.print(result)
|
||
except Exception as e:
|
||
console.print(f"[bold red]错误:[/bold red] {str(e)}")
|
||
|
||
|
||
async def test_extract_reaction_actions():
|
||
"""测试从文本提取反应步骤功能"""
|
||
console.print("\n[bold cyan]测试从文本提取反应步骤功能[/bold cyan]")
|
||
|
||
# 使用固定参数：从文本描述中提取反应步骤
|
||
fixed_reaction_text = """To a stirred solution of 7-(difluoromethylsulfonyl)-4-fluoro-indan-1-one (110 mg, 0.42 mmol) in methanol (4 mL) was added sodium borohydride (24 mg, 0.62 mmol). The reaction mixture was stirred at ambient temperature for 1 hour."""
|
||
console.print(f"固定反应文本: {fixed_reaction_text}")
|
||
|
||
try:
|
||
result = await extract_reaction_actions(fixed_reaction_text)
|
||
console.print("[green]结果:[/green]")
|
||
console.print(result)
|
||
except Exception as e:
|
||
console.print(f"[bold red]错误:[/bold red] {str(e)}")
|
||
|
||
|
||
async def test_all():
|
||
"""测试所有RXN工具函数"""
|
||
console.print("[bold magenta]===== 开始测试RXN工具函数 =====[/bold magenta]\n")
|
||
|
||
# 测试各个功能
|
||
await test_predict_reaction_outcome()
|
||
await test_predict_reaction_topn()
|
||
await test_predict_retrosynthesis()
|
||
await test_predict_biocatalytic_retrosynthesis()
|
||
await test_predict_reaction_properties()
|
||
await test_extract_reaction_actions()
|
||
|
||
console.print("\n[bold magenta]===== RXN工具函数测试完成 =====[/bold magenta]")
|
||
|
||
|
||
def get_rxn_tool_questions():
|
||
"""
|
||
获取为RXN工具函数生成的问题列表
|
||
这些问题设计为能够引导大模型调用相应的工具函数
|
||
|
||
Returns:
|
||
包含工具名称和对应问题的字典
|
||
"""
|
||
questions = {
|
||
"predict_reaction_outcome": [
|
||
"如果我将溴和蒽混合在一起，会形成什么产物？",
|
||
"乙酸和乙醇反应会生成什么？",
|
||
"预测一下丙烯醛和甲胺反应的结果"
|
||
],
|
||
|
||
|
||
"predict_reaction_topn": [
|
||
"丙烯醛和甲胺反应可能生成哪几种主要产物？",
|
||
"预测溴和蒽反应可能的前3个产物",
|
||
"乙酸和乙醇反应可能有哪些不同的结果？请给出最可能的几种产物"
|
||
],
|
||
|
||
"predict_retrosynthesis": [
|
||
"如何合成阿司匹林？请给出可能的合成路线",
|
||
"对于分子CC(=O)OC1=CC=CC=C1C(=O)O，有哪些可能的合成路径？",
|
||
"请分析一下布洛芬的可能合成路线"
|
||
],
|
||
|
||
"predict_biocatalytic_retrosynthesis": [
|
||
"有没有可能用酶催化合成OC1C(O)C=C(Br)C=C1这个分子？",
|
||
"请提供一种使用生物催化方法合成对羟基苯甲醇的路线",
|
||
"我想用酶催化方法合成一些复杂分子，能否分析一下OC1C(O)C=C(Br)C=C1的可能合成路径？"
|
||
],
|
||
|
||
"predict_reaction_properties": [
|
||
"在这个反应中：CC(C)S.CN(C)C=O.Fc1cccnc1F.O=C([O-])[O-].[K+].[K+]>>CC(C)Sc1ncccc1F，原子是如何映射的？",
|
||
"这个反应的产率可能是多少：Clc1ccccn1.Cc1ccc(N)cc1>>Cc1ccc(Nc2ccccn2)cc1",
|
||
"能分析一下这个反应中原子的去向吗：CC(=O)O.CCO>>CC(=O)OCC"
|
||
],
|
||
|
||
"extract_reaction_actions": [
|
||
"能否将这段实验描述转换为结构化的反应步骤？'To a stirred solution of 7-(difluoromethylsulfonyl)-4-fluoro-indan-1-one (110 mg, 0.42 mmol) in methanol (4 mL) was added sodium borohydride (24 mg, 0.62 mmol). The reaction mixture was stirred at ambient temperature for 1 hour.'",
|
||
"请从这段文本中提取出具体的反应操作步骤：'A solution of benzoic acid (1.0 g, 8.2 mmol) in thionyl chloride (10 mL) was heated under reflux for 2 hours. The excess thionyl chloride was removed under reduced pressure to give benzoyl chloride as a colorless liquid.'",
|
||
"帮我解析这个实验步骤，提取出关键操作：'The aldehyde (5 mmol) was dissolved in methanol (20 mL) and sodium borohydride (7.5 mmol) was added portionwise at 0°C. The mixture was allowed to warm to room temperature and stirred for 3 hours.'"
|
||
]
|
||
}
|
||
|
||
return questions
|
||
|
||
|
||
def update_agent_test_questions():
|
||
"""
|
||
更新agent_test.py中的工具问题字典，添加RXN工具函数的问题
|
||
"""
|
||
try:
|
||
# 获取agent_test.py文件路径
|
||
agent_test_path = Path('/home/ubuntu/sas0/lzy/multi_mcp_server/test_tools/agent_test.py')
|
||
|
||
# 读取文件内容
|
||
with open(agent_test_path, 'r') as f:
|
||
content = f.read()
|
||
|
||
# 获取RXN工具函数的问题
|
||
rxn_questions = get_rxn_tool_questions()
|
||
|
||
# 检查文件中是否已包含RXN工具函数的问题
|
||
rxn_tools_exist = any(tool in content for tool in rxn_questions.keys())
|
||
|
||
if not rxn_tools_exist:
|
||
# 找到questions字典的结束位置
|
||
dict_end_pos = content.find(' return questions')
|
||
|
||
if dict_end_pos != -1:
|
||
# 构建要插入的RXN工具函数问题
|
||
rxn_questions_str = ""
|
||
for tool_name, questions_list in rxn_questions.items():
|
||
rxn_questions_str += f'\n "{tool_name}": [\n'
|
||
for q in questions_list:
|
||
rxn_questions_str += f' "{q}",\n'
|
||
rxn_questions_str += ' ],'
|
||
|
||
# 在字典结束前插入RXN工具函数问题
|
||
new_content = content[:dict_end_pos] + rxn_questions_str + content[dict_end_pos:]
|
||
|
||
# 写回文件
|
||
with open(agent_test_path, 'w') as f:
|
||
f.write(new_content)
|
||
|
||
console.print("[green]成功更新agent_test.py，添加了RXN工具函数的测试问题[/green]")
|
||
else:
|
||
console.print("[yellow]无法找到questions字典的结束位置，未更新agent_test.py[/yellow]")
|
||
else:
|
||
console.print("[yellow]agent_test.py中已包含RXN工具函数的问题，无需更新[/yellow]")
|
||
|
||
except Exception as e:
|
||
console.print(f"[bold red]更新agent_test.py时出错:[/bold red] {str(e)}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 运行所有测试
|
||
# asyncio.run(test_all())
|
||
|
||
# # 更新agent_test.py中的工具问题
|
||
# update_agent_test_questions()
|
||
|
||
api_key = 'apk-8928522a146c2503f30b16d9909222d7583f412ee8f1049f08d32a089ba88d34'
|
||
from rxn4chemistry import RXN4ChemistryWrapper
|
||
|
||
rxn4chemistry_wrapper = RXN4ChemistryWrapper(api_key=api_key)
|
||
rxn4chemistry_wrapper.create_project('test_wrapper')
|
||
response = rxn4chemistry_wrapper.predict_automatic_retrosynthesis(
|
||
'Brc1c2ccccc2c(Br)c2ccccc12')
|
||
results = rxn4chemistry_wrapper.get_predict_automatic_retrosynthesis_results(response['prediction_id'])
|
||
print(results['status'])
|
||
# NOTE: upon 'SUCCESS' you can inspect the predicted retrosynthetic paths.
|
||
print(results['retrosynthetic_paths'][0])
|