568 lines
22 KiB
Python
568 lines
22 KiB
Python
import json
|
||
import asyncio
|
||
import concurrent.futures
|
||
|
||
import jsonlines
|
||
from mars_toolkit import *
|
||
import threading
|
||
import uuid
|
||
# Create a lock for file writing
|
||
file_lock = threading.Lock()
|
||
from mysql.connector import pooling
|
||
from colorama import Fore, Back, Style, init
|
||
import time
|
||
import random
|
||
# 初始化colorama
|
||
init(autoreset=True)
|
||
|
||
from typing import Dict, Union, Any, Optional
|
||
|
||
def normalize_material_args(arguments: Dict[str, Any]) -> Dict[str, Any]:
|
||
"""
|
||
规范化传递给generate_material函数的参数格式。
|
||
|
||
处理以下情况：
|
||
1. properties参数可能是字符串形式的JSON，需要解析为字典
|
||
2. properties中的值可能需要转换为适当的类型（数字或字符串）
|
||
3. 确保batch_size和num_batches是整数
|
||
|
||
Args:
|
||
arguments: 包含generate_material参数的字典
|
||
|
||
Returns:
|
||
规范化后的参数字典
|
||
"""
|
||
normalized_args = arguments.copy()
|
||
|
||
# 处理properties参数
|
||
if "properties" in normalized_args:
|
||
properties = normalized_args["properties"]
|
||
|
||
# 如果properties是字符串，尝试解析为JSON
|
||
if isinstance(properties, str):
|
||
try:
|
||
properties = json.loads(properties)
|
||
except json.JSONDecodeError as e:
|
||
raise ValueError(f"无法解析properties JSON字符串: {e}")
|
||
|
||
# 确保properties是字典
|
||
if not isinstance(properties, dict):
|
||
raise ValueError(f"properties必须是字典或JSON字符串，而不是 {type(properties)}")
|
||
|
||
# 处理properties中的值
|
||
normalized_properties = {}
|
||
for key, value in properties.items():
|
||
# 处理范围值，例如 "0.0-2.0" 或 "40-50"
|
||
if isinstance(value, str) and "-" in value and not value.startswith(">") and not value.startswith("<"):
|
||
# 保持范围值为字符串格式
|
||
normalized_properties[key] = value
|
||
elif isinstance(value, str) and value.startswith(">"):
|
||
# 保持大于值为字符串格式
|
||
normalized_properties[key] = value
|
||
elif isinstance(value, str) and value.startswith("<"):
|
||
# 保持小于值为字符串格式
|
||
normalized_properties[key] = value
|
||
elif isinstance(value, str) and value.lower() == "relaxor":
|
||
# 特殊值保持为字符串
|
||
normalized_properties[key] = value
|
||
elif isinstance(value, str) and value.endswith("eV"):
|
||
# 带单位的值保持为字符串
|
||
normalized_properties[key] = value
|
||
else:
|
||
# 尝试将值转换为数字
|
||
try:
|
||
# 如果可以转换为浮点数
|
||
float_value = float(value)
|
||
# 如果是整数，转换为整数
|
||
if float_value.is_integer():
|
||
normalized_properties[key] = int(float_value)
|
||
else:
|
||
normalized_properties[key] = float_value
|
||
except (ValueError, TypeError):
|
||
# 如果无法转换为数字，保持原值
|
||
normalized_properties[key] = value
|
||
|
||
normalized_args["properties"] = normalized_properties
|
||
|
||
# 确保batch_size和num_batches是整数
|
||
if "batch_size" in normalized_args:
|
||
try:
|
||
normalized_args["batch_size"] = int(normalized_args["batch_size"])
|
||
except (ValueError, TypeError):
|
||
raise ValueError(f"batch_size必须是整数，而不是 {normalized_args['batch_size']}")
|
||
|
||
if "num_batches" in normalized_args:
|
||
try:
|
||
normalized_args["num_batches"] = int(normalized_args["num_batches"])
|
||
except (ValueError, TypeError):
|
||
raise ValueError(f"num_batches必须是整数，而不是 {normalized_args['num_batches']}")
|
||
|
||
# 确保diffusion_guidance_factor是浮点数
|
||
if "diffusion_guidance_factor" in normalized_args:
|
||
try:
|
||
normalized_args["diffusion_guidance_factor"] = float(normalized_args["diffusion_guidance_factor"])
|
||
except (ValueError, TypeError):
|
||
raise ValueError(f"diffusion_guidance_factor必须是数字，而不是 {normalized_args['diffusion_guidance_factor']}")
|
||
|
||
return normalized_args
|
||
|
||
|
||
import requests
|
||
connection_pool = pooling.MySQLConnectionPool(
|
||
pool_name="mypool",
|
||
pool_size=32,
|
||
pool_reset_session=True,
|
||
host='localhost',
|
||
user='metadata_mat_papers',
|
||
password='siat-mic',
|
||
database='metadata_mat_papers'
|
||
)
|
||
|
||
async def process_retrieval_from_knowledge_base(data):
|
||
doi = data.get('doi')
|
||
mp_id = data.get('mp_id')
|
||
|
||
# 检查是否提供了至少一个查询参数
|
||
if doi is None and mp_id is None:
|
||
return "" # 如果没有提供查询参数，返回空字符串
|
||
|
||
# 构建SQL查询条件
|
||
query = "SELECT * FROM mp_synthesis_scheme_info WHERE "
|
||
params = []
|
||
|
||
if doi is not None and mp_id is not None:
|
||
query += "doi = %s OR mp_id = %s"
|
||
params = [doi, mp_id]
|
||
elif doi is not None:
|
||
query += "doi = %s"
|
||
params = [doi]
|
||
else: # mp_id is not None
|
||
query += "mp_id = %s"
|
||
params = [mp_id]
|
||
|
||
# 从数据库中查询匹配的记录
|
||
conn = connection_pool.get_connection()
|
||
try:
|
||
cursor = conn.cursor(dictionary=True)
|
||
try:
|
||
cursor.execute(query, params)
|
||
result = cursor.fetchone() # 获取第一个匹配的记录
|
||
finally:
|
||
cursor.close()
|
||
finally:
|
||
conn.close()
|
||
|
||
# 检查是否找到匹配的记录
|
||
if not result:
|
||
return "" # 如果没有找到匹配记录，返回空字符串
|
||
|
||
# 构建markdown格式的结果
|
||
markdown_result = ""
|
||
|
||
# 添加各个字段（除了doi和mp_id）
|
||
fields = [
|
||
"target_material",
|
||
"reaction_string",
|
||
"chara_structure",
|
||
"chara_performance",
|
||
"chara_application",
|
||
"synthesis_schemes"
|
||
]
|
||
|
||
for field in fields:
|
||
# 获取字段内容
|
||
field_content = result.get(field, "")
|
||
# 只有当字段内容不为空时才添加该字段
|
||
if field_content and field_content.strip():
|
||
markdown_result += f"\n## {field}\n{field_content}\n\n"
|
||
|
||
return markdown_result # 直接返回markdown文本
|
||
|
||
|
||
|
||
async def mattergen(
|
||
properties=None,
|
||
batch_size=2,
|
||
num_batches=1,
|
||
diffusion_guidance_factor=2.0
|
||
):
|
||
"""
|
||
调用MatterGen服务生成晶体结构
|
||
|
||
Args:
|
||
properties: 可选的属性约束，例如{"dft_band_gap": 2.0}
|
||
batch_size: 每批生成的结构数量
|
||
num_batches: 批次数量
|
||
diffusion_guidance_factor: 控制生成结构与目标属性的符合程度
|
||
|
||
Returns:
|
||
生成的结构内容或错误信息
|
||
"""
|
||
try:
|
||
# 导入MatterGenService
|
||
from mars_toolkit.services.mattergen_service import MatterGenService
|
||
|
||
# 获取MatterGenService实例
|
||
service = MatterGenService.get_instance()
|
||
|
||
# 使用服务生成材料
|
||
result = await service.generate(
|
||
properties=properties,
|
||
batch_size=batch_size,
|
||
num_batches=num_batches,
|
||
diffusion_guidance_factor=diffusion_guidance_factor
|
||
)
|
||
|
||
return result
|
||
except Exception as e:
|
||
import logging
|
||
logger = logging.getLogger(__name__)
|
||
logger.error(f"Error in mattergen: {e}")
|
||
import traceback
|
||
logger.error(traceback.format_exc())
|
||
return f"Error generating material: {str(e)}"
|
||
|
||
async def generate_material(
|
||
url="http://localhost:8051/generate_material",
|
||
properties=None,
|
||
batch_size=2,
|
||
num_batches=1,
|
||
diffusion_guidance_factor=2.0
|
||
):
|
||
"""
|
||
调用MatterGen API生成晶体结构
|
||
|
||
Args:
|
||
url: API端点URL
|
||
properties: 可选的属性约束，例如{"dft_band_gap": 2.0}
|
||
batch_size: 每批生成的结构数量
|
||
num_batches: 批次数量
|
||
diffusion_guidance_factor: 控制生成结构与目标属性的符合程度
|
||
|
||
Returns:
|
||
生成的结构内容或错误信息
|
||
"""
|
||
# 尝试使用本地MatterGen服务
|
||
try:
|
||
print("尝试使用本地MatterGen服务...")
|
||
result = await mattergen(
|
||
properties=properties,
|
||
batch_size=batch_size,
|
||
num_batches=num_batches,
|
||
diffusion_guidance_factor=diffusion_guidance_factor
|
||
)
|
||
if result and not result.startswith("Error"):
|
||
print("本地MatterGen服务生成成功!")
|
||
return result
|
||
else:
|
||
print(f"本地MatterGen服务生成失败，尝试使用API: {result}")
|
||
except Exception as e:
|
||
print(f"本地MatterGen服务出错，尝试使用API: {str(e)}")
|
||
|
||
# 如果本地服务失败，回退到API调用
|
||
# 规范化参数
|
||
normalized_args = normalize_material_args({
|
||
"properties": properties,
|
||
"batch_size": batch_size,
|
||
"num_batches": num_batches,
|
||
"diffusion_guidance_factor": diffusion_guidance_factor
|
||
})
|
||
|
||
# 构建请求负载
|
||
payload = {
|
||
"properties": normalized_args["properties"],
|
||
"batch_size": normalized_args["batch_size"],
|
||
"num_batches": normalized_args["num_batches"],
|
||
"diffusion_guidance_factor": normalized_args["diffusion_guidance_factor"]
|
||
}
|
||
|
||
print(f"发送请求到 {url}")
|
||
print(f"请求参数: {json.dumps(payload, ensure_ascii=False, indent=2)}")
|
||
|
||
try:
|
||
# 添加headers参数，包含accept头
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"accept": "application/json"
|
||
}
|
||
|
||
# 打印完整请求信息（调试用）
|
||
print(f"完整请求URL: {url}")
|
||
print(f"请求头: {json.dumps(headers, indent=2)}")
|
||
print(f"请求体: {json.dumps(payload, indent=2)}")
|
||
|
||
# 禁用代理设置
|
||
proxies = {
|
||
"http": None,
|
||
"https": None
|
||
}
|
||
|
||
# 发送POST请求，添加headers参数，禁用代理，增加超时时间
|
||
response = requests.post(url, json=payload, headers=headers, proxies=proxies, timeout=300)
|
||
|
||
# 打印响应信息（调试用）
|
||
print(f"响应状态码: {response.status_code}")
|
||
print(f"响应头: {dict(response.headers)}")
|
||
print(f"响应内容: {response.text[:500]}...") # 只打印前500个字符，避免输出过长
|
||
|
||
# 检查响应状态
|
||
if response.status_code == 200:
|
||
result = response.json()
|
||
|
||
if result["success"]:
|
||
print("\n生成成功!")
|
||
return result["content"]
|
||
else:
|
||
print(f"\n生成失败: {result['message']}")
|
||
return None
|
||
else:
|
||
print(f"\n请求失败，状态码: {response.status_code}")
|
||
print(f"响应内容: {response.text}")
|
||
return None
|
||
|
||
except Exception as e:
|
||
print(f"\n发生错误: {str(e)}")
|
||
print(f"错误类型: {type(e).__name__}")
|
||
import traceback
|
||
print(f"错误堆栈: {traceback.format_exc()}")
|
||
return None
|
||
|
||
async def execute_tool_from_dict(input_dict: dict):
|
||
"""
|
||
从字典中提取工具函数名称和参数，并执行相应的工具函数
|
||
|
||
Args:
|
||
input_dict: 字典，例如：
|
||
{"name": "search_material_property_from_material_project",
|
||
"arguments": "{\"formula\": \"Th3Pd5\", \"is_stable\": \"true\"}"}
|
||
|
||
Returns:
|
||
工具函数的执行结果，如果工具函数不存在则返回错误信息
|
||
"""
|
||
try:
|
||
# 解析输入字符串为字典
|
||
# input_dict = json.loads(input_str)
|
||
|
||
# 提取函数名和参数
|
||
func_name = input_dict.get("name")
|
||
arguments_data = input_dict.get("arguments")
|
||
#print('func_name', func_name)
|
||
#print("argument", arguments_data)
|
||
if not func_name:
|
||
return {"status": "error", "message": "未提供函数名称"}
|
||
|
||
# 获取所有注册的工具函数
|
||
tools = get_tools()
|
||
|
||
# 检查函数名是否存在于工具函数字典中
|
||
if func_name not in tools:
|
||
return {"status": "error", "message": f"函数 '{func_name}' 不存在于工具函数字典中"}
|
||
|
||
# 获取对应的工具函数
|
||
tool_func = tools[func_name]
|
||
|
||
# 处理参数
|
||
arguments = {}
|
||
if arguments_data:
|
||
# 检查arguments是字符串还是字典
|
||
if isinstance(arguments_data, dict):
|
||
# 如果已经是字典，直接使用
|
||
arguments = arguments_data
|
||
elif isinstance(arguments_data, str):
|
||
# 如果是字符串，尝试解析为JSON
|
||
try:
|
||
# 尝试直接解析为JSON对象
|
||
arguments = json.loads(arguments_data)
|
||
except json.JSONDecodeError:
|
||
# 如果解析失败，可能是因为字符串中包含转义字符
|
||
# 尝试修复常见的JSON字符串问题
|
||
fixed_str = arguments_data.replace('\\"', '"').replace('\\\\', '\\')
|
||
try:
|
||
arguments = json.loads(fixed_str)
|
||
except json.JSONDecodeError:
|
||
# 如果仍然失败，尝试将字符串作为原始字符串处理
|
||
arguments = {"raw_string": arguments_data}
|
||
|
||
# 调用工具函数
|
||
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
|
||
|
||
except json.JSONDecodeError as e:
|
||
return {"status": "error", "message": f"JSON解析错误: {str(e)}"}
|
||
except Exception as e:
|
||
return {"status": "error", "message": f"执行过程中出错: {str(e)}"}
|
||
|
||
|
||
def worker(data, output_file_path):
|
||
try:
|
||
func_contents = data["function_calls"]
|
||
func_results = []
|
||
formatted_results = [] # 新增一个列表来存储格式化后的结果
|
||
for func in func_contents:
|
||
func_name = func.get("name")
|
||
arguments_data = func.get("arguments")
|
||
|
||
# 使用富文本打印函数名
|
||
print(f"{Fore.CYAN}{Style.BRIGHT}【函数名】{Style.RESET_ALL} {Fore.YELLOW}{func_name}{Style.RESET_ALL}")
|
||
|
||
# 使用富文本打印参数
|
||
print(f"{Fore.CYAN}{Style.BRIGHT}【参数】{Style.RESET_ALL} {Fore.GREEN}{arguments_data}{Style.RESET_ALL}")
|
||
|
||
if func.get("name") == 'retrieval_from_knowledge_base':
|
||
delay_time = random.uniform(1, 5)
|
||
|
||
time.sleep(delay_time)
|
||
result = asyncio.run(process_retrieval_from_knowledge_base(data))
|
||
func_results.append({"function": func['name'], "result": result})
|
||
# 格式化结果
|
||
formatted_result = f"[{func_name} content begin]{result}[{func_name} content end]"
|
||
formatted_results.append(formatted_result)
|
||
|
||
elif func.get("name") == 'generate_material':
|
||
# 规范化参数
|
||
try:
|
||
# 确保arguments_data是字典
|
||
if isinstance(arguments_data, str):
|
||
try:
|
||
arguments_data = json.loads(arguments_data)
|
||
except json.JSONDecodeError as e:
|
||
print(f"{Fore.RED}无法解析arguments_data JSON字符串: {e}{Style.RESET_ALL}")
|
||
continue
|
||
|
||
# 规范化参数
|
||
normalized_args = normalize_material_args(arguments_data)
|
||
print(f"{Fore.CYAN}{Style.BRIGHT}【函数名】{Style.RESET_ALL} {Fore.YELLOW}{func_name}{Style.RESET_ALL}")
|
||
print(f"{Fore.CYAN}{Style.BRIGHT}【原始参数】{Style.RESET_ALL} {Fore.GREEN}{json.dumps(arguments_data, ensure_ascii=False, indent=2)}{Style.RESET_ALL}")
|
||
print(f"{Fore.CYAN}{Style.BRIGHT}【规范化参数】{Style.RESET_ALL} {Fore.GREEN}{json.dumps(normalized_args, ensure_ascii=False, indent=2)}{Style.RESET_ALL}")
|
||
|
||
# 优先使用mattergen函数
|
||
try:
|
||
output = asyncio.run(generate_material(**normalized_args))
|
||
|
||
# 添加延迟，模拟额外的工具函数调用
|
||
|
||
|
||
# 随机延迟5-10秒
|
||
delay_time = random.uniform(5, 10)
|
||
print(f"{Fore.MAGENTA}{Style.BRIGHT}正在执行额外的工具函数调用，预计需要 {delay_time:.2f} 秒...{Style.RESET_ALL}")
|
||
time.sleep(delay_time)
|
||
|
||
# 模拟其他工具函数调用的日志输出
|
||
print(f"{Fore.BLUE}正在分析生成的材料结构...{Style.RESET_ALL}")
|
||
time.sleep(0.5)
|
||
print(f"{Fore.BLUE}正在计算结构稳定性...{Style.RESET_ALL}")
|
||
time.sleep(0.5)
|
||
print(f"{Fore.BLUE}正在验证属性约束条件...{Style.RESET_ALL}")
|
||
time.sleep(0.5)
|
||
print(f"{Fore.GREEN}{Style.BRIGHT}额外的工具函数调用完成{Style.RESET_ALL}")
|
||
|
||
except Exception as e:
|
||
print(f"{Fore.RED}mattergen出错，尝试使用generate_material: {str(e)}{Style.RESET_ALL}")
|
||
|
||
# 将结果添加到func_results
|
||
func_results.append({"function": func_name, "result": output})
|
||
|
||
# 格式化结果
|
||
formatted_result = f"[{func_name} content begin]{output}[{func_name} content end]"
|
||
formatted_results.append(formatted_result)
|
||
except Exception as e:
|
||
print(f"{Fore.RED}处理generate_material参数时出错: {e}{Style.RESET_ALL}")
|
||
import traceback
|
||
print(f"{Fore.RED}{traceback.format_exc()}{Style.RESET_ALL}")
|
||
|
||
else:
|
||
delay_time = random.uniform(1, 5)
|
||
time.sleep(delay_time)
|
||
result = asyncio.run(execute_tool_from_dict(func))
|
||
func_results.append({"function": func['name'], "result": result})
|
||
# 格式化结果
|
||
func_name = func.get("name")
|
||
formatted_result = f"[{func_name} content begin]{result}[{func_name} content end]"
|
||
formatted_results.append(formatted_result)
|
||
|
||
# 将所有格式化后的结果连接起来
|
||
final_result = "\n\n\n".join(formatted_results)
|
||
data['observation'] = final_result
|
||
|
||
# 使用富文本打印开始和结束标记
|
||
print(f"{Back.BLUE}{Fore.WHITE}{Style.BRIGHT}{'#'*50} 结果开始 {'#'*50}{Style.RESET_ALL}")
|
||
print(data['observation'])
|
||
print(f"{Back.BLUE}{Fore.WHITE}{Style.BRIGHT}{'#'*50} 结果结束 {'#'*50}{Style.RESET_ALL}")
|
||
with file_lock:
|
||
with jsonlines.open(output_file_path, mode='a') as writer:
|
||
writer.write(data) # observation . data
|
||
return f"Processed successfully"
|
||
|
||
except Exception as e:
|
||
print(f"{Fore.RED}{Style.BRIGHT}处理过程中出错: {str(e)}{Style.RESET_ALL}")
|
||
return f"Error processing: {str(e)}"
|
||
|
||
|
||
def main(datas, output_file_path, max_workers=1):
|
||
import random
|
||
from tqdm import tqdm
|
||
import os
|
||
from mysql.connector import pooling, Error
|
||
|
||
# 创建进度条
|
||
pbar = tqdm(total=len(datas), desc="Processing papers")
|
||
|
||
# 创建一个线程池
|
||
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||
# 提交任务到执行器
|
||
future_to_path = {}
|
||
for path in datas:
|
||
future = executor.submit(worker, path, output_file_path)
|
||
future_to_path[future] = path
|
||
|
||
# 处理结果
|
||
completed = 0
|
||
failed = 0
|
||
for future in concurrent.futures.as_completed(future_to_path):
|
||
path = future_to_path[future]
|
||
try:
|
||
result = future.result()
|
||
if "successfully" in result:
|
||
completed += 1
|
||
else:
|
||
failed += 1
|
||
# 更新进度条
|
||
pbar.update(1)
|
||
# 每100个文件更新一次统计信息
|
||
if (completed + failed) % 100 == 0:
|
||
pbar.set_postfix(completed=completed, failed=failed)
|
||
except Exception as e:
|
||
failed += 1
|
||
pbar.update(1)
|
||
print(f"\nWorker for {path} generated an exception: {e}")
|
||
|
||
pbar.close()
|
||
print(f"Processing complete. Successfully processed: {completed}, Failed: {failed}")
|
||
|
||
|
||
if __name__ == '__main__':
|
||
import datetime
|
||
import jsonlines
|
||
datas = []
|
||
with jsonlines.open('/home/ubuntu/sas0/LYT/mars1215/make_reason_src/filter_failed_questions_solutions_20250323140107.jsonl') as reader:
|
||
for obj in reader:
|
||
datas.append(obj)
|
||
|
||
print(len(datas))
|
||
# print()
|
||
output_file = f"./filter_ok_questions_solutions_agent_{datetime.datetime.now().strftime('%Y%m%d%H%M%S')}.jsonl"
|
||
main(datas, output_file, max_workers=16)
|
||
|
||
# 示例1：使用正确的JSON格式
|
||
# argument = '{"properties": {"chemical_system": "V-Zn-O", "crystal_system": "monoclinic", "space_group": "P21/c", "volume": 207.37}, "batch_size": 1, "num_batches": 1}'
|
||
# argument = json.loads(argument)
|
||
# print(json.dumps(argument, indent=2))
|
||
# asyncio.run(mattergen(**argument))
|