import glob import json import asyncio import logging import datetime import os from multiprocessing import Process, Manager from typing import Dict, Any, List, Optional from mp_api.client import MPRester from pymatgen.core import Structure from pymatgen.symmetry.analyzer import SpacegroupAnalyzer from pymatgen.io.cif import CifWriter from mars_toolkit.core.llm_tools import llm_tool from mars_toolkit.core.config import config from mars_toolkit.core.cif_utils import read_cif_txt_file, extract_cif_info, remove_symmetry_equiv_xyz logger = logging.getLogger(__name__) def parse_bool(param: str) -> bool | None: """ Parse a string parameter into a boolean value. Args: param: String parameter to parse (e.g., "true", "false") Returns: Boolean value if param is not empty, None otherwise """ if not param: return None return param.lower() == 'true' def parse_list(param: str) -> List[str] | None: """ Parse a comma-separated string into a list of strings. Args: param: Comma-separated string (e.g., "Li,Fe,O") Returns: List of strings if param is not empty, None otherwise """ if not param: return None return param.split(',') def parse_tuple(param: str) -> tuple[float, float] | None: """ Parse a comma-separated string into a tuple of two float values. Used for range parameters like band_gap, density, etc. Args: param: Comma-separated string of two numbers (e.g., "0,3.5") Returns: Tuple of two float values if param is valid, None otherwise """ if not param: return None try: values = param.split(',') return (float(values[0]), float(values[1])) except (ValueError, IndexError): return None def parse_search_parameters(query_params: Dict[str, str]) -> Dict[str, Any]: """ Parse search parameters from query parameters. Converts string query parameters into appropriate types for the Materials Project API. """ return { 'band_gap': parse_tuple(query_params.get('band_gap')), 'chemsys': parse_list(query_params.get('chemsys')), 'crystal_system': parse_list(query_params.get('crystal_system')), 'density': parse_tuple(query_params.get('density')), 'formation_energy': parse_tuple(query_params.get('formation_energy')), 'elements': parse_list(query_params.get('elements')), 'exclude_elements': parse_list(query_params.get('exclude_elements')), 'formula': parse_list(query_params.get('formula')), 'is_gap_direct': parse_bool(query_params.get('is_gap_direct')), 'is_metal': parse_bool(query_params.get('is_metal')), 'is_stable': parse_bool(query_params.get('is_stable')), 'magnetic_ordering': query_params.get('magnetic_ordering'), 'material_ids': parse_list(query_params.get('material_ids')), 'total_energy': parse_tuple(query_params.get('total_energy')), 'num_elements': parse_tuple(query_params.get('num_elements')), 'volume': parse_tuple(query_params.get('volume')), 'chunk_size': int(query_params.get('chunk_size', '5')) } def process_search_results(docs: List[Dict[str, Any]]) -> List[Dict[str, Any]] | str: """ Process search results from the Materials Project API. Extracts relevant fields from each document and formats them into a consistent structure. Returns: List of processed documents or error message string if an exception occurs """ try: fields = [ 'formula_pretty', 'nsites', 'nelements', 'material_id', 'chemsys', 'volume', 'density', 'density_atomic', 'cbm', 'vbm', 'band_gap', 'is_gap_direct', 'is_stable', 'formation_energy_per_atom', 'energy_above_hull', 'is_metal', 'total_magnetization', 'efermi', 'is_magnetic', 'ordering', 'bulk_modulus', 'shear_modulus', 'universal_anisotropy', 'theoretical' ] res = [] for doc in docs: try: new_docs = {} for field_name in fields: new_docs[field_name] = doc.get(field_name, '') res.append(new_docs) except Exception as e: logger.warning(f"Error processing document: {str(e)}") continue return res except Exception as e: error_msg = f"Error in process_search_results: {str(e)}" logger.error(error_msg) return error_msg def _search_worker(queue, api_key, **kwargs): """ Worker function for executing Materials Project API searches. Runs in a separate process to perform the actual API call and puts results in the queue. Args: queue: Multiprocessing queue for returning results api_key: Materials Project API key **kwargs: Search parameters to pass to the API """ try: import os import traceback os.environ['HTTP_PROXY'] = config.HTTP_PROXY or '' os.environ['HTTPS_PROXY'] = config.HTTPS_PROXY or '' # 初始化 MPRester 客户端 with MPRester(api_key) as mpr: result = mpr.materials.summary.search(**kwargs) # 检查结果 if result: # 尝试使用更安全的方式处理结果 processed_results = [] for doc in result: try: # 尝试使用 model_dump 方法 processed_doc = doc.model_dump() processed_results.append(processed_doc) except AttributeError: # 如果没有 model_dump 方法，尝试使用 dict 方法 try: processed_doc = doc.dict() processed_results.append(processed_doc) except AttributeError: # 如果没有 dict 方法，尝试直接转换为字典 if hasattr(doc, "__dict__"): processed_doc = doc.__dict__ # 移除可能导致序列化问题的特殊属性 if "_sa_instance_state" in processed_doc: del processed_doc["_sa_instance_state"] processed_results.append(processed_doc) else: # 最后的尝试，直接使用 doc processed_results.append(doc) queue.put(processed_results) else: queue.put([]) except Exception as e: queue.put(e) async def execute_search(search_args: Dict[str, Any], timeout: int = 120) -> List[Dict[str, Any]] | str: """ Execute a search against the Materials Project API. Runs the search in a separate process to handle potential timeouts and returns the results. Args: search_args: Dictionary of search parameters timeout: Maximum time in seconds to wait for the search to complete Returns: List of document dictionaries from the search results or error message string if an exception occurs """ # 确保 formula 参数是列表类型 if 'formula' in search_args and isinstance(search_args['formula'], str): search_args['formula'] = [search_args['formula']] manager = Manager() queue = manager.Queue() try: p = Process(target=_search_worker, args=(queue, config.MP_API_KEY), kwargs=search_args) p.start() p.join(timeout=timeout) if p.is_alive(): logger.warning(f"Terminating worker process {p.pid} due to timeout") p.terminate() p.join() error_msg = f"Request timed out after {timeout} seconds" return error_msg try: result = queue.get(timeout=timeout) if isinstance(result, Exception): logger.error(f"Error in search worker: {str(result)}") if hasattr(result, "__traceback__"): import traceback tb_str = ''.join(traceback.format_exception(None, result, result.__traceback__)) return f"Error in search worker: {str(result)}" return result except queue.Empty: error_msg = "Failed to retrieve data from queue (timeout)" logger.error(error_msg) return error_msg except Exception as e: error_msg = f"Error in execute_search: {str(e)}" logger.error(error_msg) return error_msg @llm_tool(name="search_material_property_from_material_project", description="Search materials in Materials Project database by formula and properties") async def search_material_property_from_material_project( formula: str | list[str], chemsys: Optional[str | list[str] | None] = None, crystal_system: Optional[str | list[str] | None] = None, is_gap_direct: Optional[bool | None] = None, is_stable: Optional[bool | None] = None, ) -> str: """ Search materials in Materials Project database. Args: formula: Chemical formula(s) (e.g., "Fe2O3" or ["ABO3", "Si*"]) chemsys: Chemical system(s) (e.g., "Li-Fe-O") crystal_system: Crystal system(s) (e.g., "Cubic") is_gap_direct: Filter for direct band gap materials is_stable: Filter for thermodynamically stable materials Returns: JSON formatted material properties data """ # 验证晶系参数 VALID_CRYSTAL_SYSTEMS = ['Triclinic', 'Monoclinic', 'Orthorhombic', 'Tetragonal', 'Trigonal', 'Hexagonal', 'Cubic'] # 验证晶系参数是否有效 if crystal_system is not None: if isinstance(crystal_system, str): if crystal_system not in VALID_CRYSTAL_SYSTEMS: return "Input should be 'Triclinic', 'Monoclinic', 'Orthorhombic', 'Tetragonal', 'Trigonal', 'Hexagonal' or 'Cubic'" elif isinstance(crystal_system, list): for cs in crystal_system: if cs not in VALID_CRYSTAL_SYSTEMS: return "Input should be 'Triclinic', 'Monoclinic', 'Orthorhombic', 'Tetragonal', 'Trigonal', 'Hexagonal' or 'Cubic'" # 确保 formula 是列表类型 if isinstance(formula, str): formula = [formula] params = { "chemsys": chemsys, "crystal_system": crystal_system, "formula": formula, "is_gap_direct": is_gap_direct, "is_stable": is_stable, "chunk_size": 5, } # Filter out None values params = {k: v for k, v in params.items() if v is not None} mp_id_list = await get_mpid_from_formula(formula=formula) try: res=[] for mp_id in mp_id_list: crystal_props = extract_cif_info(config.LOCAL_MP_ROOT+f"/Props/{mp_id}.json", ['all_fields']) res.append(crystal_props) if len(res) == 0: return "No results found, please try again." # Format response with top results try: # 创建包含索引的JSON结果 formatted_results = [] for i, item in enumerate(res[:config.MP_TOPK], 1): formatted_result = f"[property {i} begin]\n" formatted_result += json.dumps(item, indent=2) formatted_result += f"\n[property {i} end]\n\n" formatted_results.append(formatted_result) # 将所有结果合并为一个字符串 res_chunk = "\n\n".join(formatted_results) res_template = f""" Here are the search results from the Materials Project database: Due to length limitations, only the top {config.MP_TOPK} results are shown below:\n {res_chunk} If you need more results, please modify your search criteria or try different query parameters. """ return res_template except Exception as format_error: logger.error(f"Error formatting results: {str(format_error)}") return str(format_error) except Exception as e: logger.error(f"Error in search_material_property_from_material_project: {str(e)}") return str(e) @llm_tool(name="get_crystal_structures_from_materials_project", description="Get symmetrized crystal structures CIF data from Materials Project database by chemical formula") async def get_crystal_structures_from_materials_project( formulas: list[str], conventional_unit_cell: bool = True, symprec: float = 0.1 ) -> str: """ Get crystal structures from Materials Project database by chemical formula and apply symmetrization. Args: formulas: List of chemical formulas (e.g., ["Fe2O3", "SiO2", "TiO2"]) conventional_unit_cell: Whether to return conventional unit cell (True) or primitive cell (False) symprec: Precision parameter for symmetrization Returns: Formatted text containing symmetrized CIF data """ result={} mp_id_list=await get_mpid_from_formula(formula=formulas) for i,mp_id in enumerate(mp_id_list): cif_file = glob.glob(config.LOCAL_MP_ROOT+f"/MPDatasets/{mp_id}.cif")[0] structure = Structure.from_file(cif_file) # 如果需要常规单元格 if conventional_unit_cell: structure = SpacegroupAnalyzer(structure).get_conventional_standard_structure() # 对结构进行对称化处理 sga = SpacegroupAnalyzer(structure, symprec=symprec) symmetrized_structure = sga.get_refined_structure() # 使用CifWriter生成CIF数据 cif_writer = CifWriter(symmetrized_structure, symprec=symprec, refine_struct=True) cif_data = str(cif_writer) # 删除CIF文件中的对称性操作部分 cif_data = remove_symmetry_equiv_xyz(cif_data) cif_data=cif_data.replace('# generated using pymatgen',"") # 生成一个唯一的键 formula = structure.composition.reduced_formula key = f"{formula}_{i}" result[key] = cif_data # 只保留前config.MP_TOPK个结果 if len(result) >= config.MP_TOPK: break try: prompt = f""" # Materials Project Symmetrized Crystal Structure Data Below are symmetrized crystal structure data for {len(result)} materials from the Materials Project database, in CIF (Crystallographic Information File) format. These structures have been analyzed and optimized for symmetry using SpacegroupAnalyzer with precision parameter symprec={symprec}.\n """ for i, (key, cif_data) in enumerate(result.items(), 1): prompt += f"[cif {i} begin]\n" prompt += cif_data prompt += f"\n[cif {i} end]\n\n" prompt += """ ## Usage Instructions 1. You can copy the above CIF data and save it as .cif files 2. Open these files with crystal structure visualization software (such as VESTA, Mercury, Avogadro, etc.) 3. These structures can be used for further material analysis, simulation, or visualization CIF files contain complete structural information of crystals, including cell parameters, atomic coordinates, symmetry, etc. Symmetrization helps identify and optimize crystal symmetry, making the structure more standardized and accurate. """ return prompt except Exception as format_error: logger.error(f"Error formatting crystal structures: {str(format_error)}") return str(format_error) @llm_tool(name="get_mpid_from_formula", description="Get material IDs (mpid) from Materials Project database by chemical formula") async def get_mpid_from_formula(formula: str) -> List[str]: """ Get material IDs (mpid) from Materials Project database by chemical formula. Returns mpids for the lowest energy structures. Args: formula: Chemical formula (e.g., "Fe2O3") Returns: List of material IDs """ os.environ['HTTP_PROXY'] = config.HTTP_PROXY or '' os.environ['HTTPS_PROXY'] = config.HTTPS_PROXY or '' id_list = [] try: with MPRester(config.MP_API_KEY) as mpr: docs = mpr.materials.summary.search(formula=formula) for doc in docs: id_list.append(doc.material_id) return id_list except Exception as e: logger.error(f"Error getting mpid from formula: {str(e)}") return []