234 lines
8.4 KiB
Python
234 lines
8.4 KiB
Python
import os
|
|
import requests
|
|
from datetime import datetime
|
|
from typing import Optional
|
|
import pytz
|
|
|
|
|
|
class Tools:
|
|
def __init__(self):
|
|
"""
|
|
Initialize the tools.
|
|
"""
|
|
self.tool_endpoint = "http://100.84.94.73:8020"
|
|
|
|
# Add your custom tools using pure Python code here, make sure to add type hints
|
|
# Use Sphinx-style docstrings to document your tools, they will be used for generating tools specifications
|
|
# Please refer to function_calling_filter_pipeline.py file from pipelines project for an example
|
|
|
|
def calculator(self, equation: str) -> str:
|
|
"""
|
|
Perform mathematical calculations on a given equation.
|
|
|
|
Args:
|
|
equation (str): A valid mathematical expression to evaluate.
|
|
Example: "2 + 3 * (4 - 1)"
|
|
|
|
Returns:
|
|
str: The calculation result in format "equation = result" or error message if invalid
|
|
"""
|
|
|
|
# Avoid using eval in production code
|
|
# https://nedbatchelder.com/blog/201206/eval_really_is_dangerous.html
|
|
try:
|
|
result = eval(equation)
|
|
return f"{equation} = {result}"
|
|
except Exception as e:
|
|
print(e)
|
|
return "Invalid equation"
|
|
|
|
def search_property_from_material_project(
|
|
self,
|
|
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 material properties from Material Project database.
|
|
|
|
Args:
|
|
formula (str | list[str]): [Required] Chemical formula(s) to search. Example: "Fe2O3" or ["ABO3", "Si*"]
|
|
chemsys (str | list[str] | None): [Optional] Chemical system(s) to search. Example: "Li-Fe-O" or ["Si-O", "Li-Fe-P"]
|
|
crystal_system (str | list[str] | None): [Optional] Crystal system(s) to filter by. Example: "Cubic" or ["Hexagonal", "Tetragonal"]
|
|
is_gap_direct (bool | None): [Optional] Filter by materials with direct band gap
|
|
is_stable (bool | None): [Optional] Filter by thermodynamic stability
|
|
|
|
Returns:
|
|
str: JSON string containing material properties or error message
|
|
"""
|
|
# Build request parameters
|
|
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}
|
|
|
|
try:
|
|
response = requests.get(
|
|
f"{self.tool_endpoint}/mp/search",
|
|
params=params,
|
|
timeout=30
|
|
)
|
|
response.raise_for_status()
|
|
return str(response.json()['data'])
|
|
|
|
except requests.exceptions.Timeout:
|
|
return "Request timed out"
|
|
except requests.exceptions.RequestException as e:
|
|
return f"Request failed: {str(e)}"
|
|
|
|
|
|
def search_from_oqmd_by_composition(self, composition: str) -> str:
|
|
"""
|
|
Search materials from OQMD database by chemical composition.
|
|
|
|
Args:
|
|
composition (str): Chemical composition string. Example: "CsPbBr3"
|
|
|
|
Returns:
|
|
str: JSON string containing material data or error message
|
|
"""
|
|
# 构建请求参数
|
|
param = {
|
|
'composition': composition
|
|
}
|
|
try:
|
|
# 发送请求到/oqmd/search路由
|
|
response = requests.get(
|
|
self.tool_endpoint + "/oqmd/search",
|
|
params=param
|
|
)
|
|
response.raise_for_status()
|
|
return str(response.json()['data'])
|
|
except requests.exceptions.RequestException as e:
|
|
return f"Error: {str(e)}"
|
|
|
|
|
|
def optimize_crystal_structure(self, markdown_block: str, input_format: str = "cif", output_format: str = "cif") -> str:
|
|
"""
|
|
Optimize crystal structure using various file formats.
|
|
|
|
Args:
|
|
markdown_block (str): Markdown block containing structure data.
|
|
Format: ```format\ncontent\n``` where format is one of ["cif", "poscar", "json", "xyz"]
|
|
input_format (str): Input file format. Must be one of ["cif", "poscar", "json", "xyz"]
|
|
output_format (str): Output file format. Must be one of ["cif", "poscar", "json", "xyz"]
|
|
|
|
Returns:
|
|
str: JSON string containing optimized structure data or error message
|
|
"""
|
|
# Validate input/output formats
|
|
valid_formats = ["cif", "poscar", "json", "xyz"]
|
|
if input_format not in valid_formats:
|
|
return f"Error: Invalid input format. Must be one of {valid_formats}"
|
|
if output_format not in valid_formats:
|
|
return f"Error: Invalid output format. Must be one of {valid_formats}"
|
|
|
|
# Define format specific validation
|
|
format_validators = {
|
|
"cif": ("```cif\n", "\n```"),
|
|
"poscar": ("```poscar\n", "\n```"),
|
|
"json": ("```json\n", "\n```"),
|
|
"xyz": ("```xyz\n", "\n```")
|
|
}
|
|
|
|
try:
|
|
# Validate markdown block format
|
|
start_tag, end_tag = format_validators[input_format]
|
|
if not markdown_block.startswith(start_tag) or not markdown_block.endswith(end_tag):
|
|
return f"Error: Invalid {input_format} markdown format. Expected format: {start_tag}...{end_tag}"
|
|
|
|
# Extract content
|
|
content = markdown_block[len(start_tag):-len(end_tag)].strip()
|
|
if not content:
|
|
return f"Error: Empty {input_format} content"
|
|
|
|
# Prepare request parameters
|
|
params = {
|
|
'input_format': input_format,
|
|
'output_format': output_format
|
|
}
|
|
|
|
# Send request to optimization service
|
|
try:
|
|
response = requests.post(
|
|
self.tool_endpoint + "/fairchem/optimize_structure",
|
|
data=content,
|
|
headers={"Content-Type": "text/plain"},
|
|
params=params,
|
|
timeout=30 # Add timeout
|
|
)
|
|
response.raise_for_status()
|
|
|
|
# 直接返回完整响应
|
|
return str(response.json()['data'])
|
|
|
|
except requests.exceptions.Timeout:
|
|
return "Error: Optimization request timed out"
|
|
except requests.exceptions.RequestException as e:
|
|
return f"Error: Optimization request failed - {str(e)}"
|
|
except ValueError as e:
|
|
return f"Error: Invalid JSON response - {str(e)}"
|
|
|
|
except Exception as e:
|
|
return f"Error: {str(e)}"
|
|
|
|
|
|
if __name__ == '__main__':
|
|
tools = Tools()
|
|
# print(tools.search_property_from_material_project(formula='CsPbBr3'))
|
|
# print(tools.search_from_material_project())
|
|
print(tools.search_from_oqmd_by_composition("CsPbBr3"))
|
|
# print(tools.optimize_crystal_structure("""```cif\n# generated using pymatgen
|
|
# data_CsPbBr3
|
|
# _symmetry_space_group_name_H-M Pnma
|
|
# _cell_length_a 8.45384704
|
|
# _cell_length_b 11.87891123
|
|
# _cell_length_c 8.10107841
|
|
# _cell_angle_alpha 90.00000000
|
|
# _cell_angle_beta 90.00000000
|
|
# _cell_angle_gamma 90.00000000
|
|
# _symmetry_Int_Tables_number 62
|
|
# _chemical_formula_structural CsPbBr3
|
|
# _chemical_formula_sum 'Cs4 Pb4 Br12'
|
|
# _cell_volume 813.53053480
|
|
# _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+1/2, -y, z+1/2'
|
|
# 4 'x+1/2, y, -z+1/2'
|
|
# 5 'x+1/2, -y+1/2, -z+1/2'
|
|
# 6 '-x+1/2, y+1/2, z+1/2'
|
|
# 7 '-x, y+1/2, -z'
|
|
# 8 'x, -y+1/2, z'
|
|
# loop_
|
|
# _atom_type_symbol
|
|
# _atom_type_oxidation_number
|
|
# Cs+ 1.0
|
|
# Pb2+ 2.0
|
|
# Br- -1.0
|
|
# 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
|
|
# Cs+ Cs0 4 0.06016347 0.75000000 0.01945212 1
|
|
# Pb2+ Pb1 4 0.00000000 0.00000000 0.50000000 1
|
|
# Br- Br2 8 0.20102217 0.03170128 0.80279189 1
|
|
# Br- Br3 4 0.00625945 0.25000000 0.44033452 1
|
|
# \n```""", input_format="cif", output_format="poscar"))
|