Files
mars_toolkit/test_client.py

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"))