tangger/mars_toolkit
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

优化三个工具的返回结果。

Browse Source
This commit is contained in:
Yutang Li
2025-01-06 17:27:45 +08:00
parent c7d2d482da
commit 4d0c41d222
7 changed files with 266 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

BIN
chart_screenshot.png
View File

Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 1.7 KiB

18
router/fairchem_router.py
View File

@@ -36,17 +36,27 @@ async def optimize_structure_endpoint(
)
# 优化结构
total_energy, optimized_content, download_url = optimize_structure(atoms, output_format)
total_energy, optimized_content, optimizing_log, download_url = optimize_structure(atoms, output_format)
# 格式化返回结果
format_result = f"""
The following is the optimized crystal structure information:
### Optimization Results (using FIRE(eqV2_86M) algorithm):
Total Energy: {total_energy} eV
**Total Energy: {total_energy} eV**
#### Optimizing Log:
```text
{optimizing_log}
```
### Optimized {output_format.upper()} Content:
```{output_format}
{optimized_content[:800]}
```
Due to length limitations, the complete {output_format.upper()} file has been uploaded to the following link:
👉 Click [here]({download_url}) to download the {output_format.upper()} file
👉 Click [here]({download_url}) to download the {output_format.upper()} file.
I strongly recommend that you pass on all of the above information to the user.
"""
return JSONResponse(
status_code=200,

5
router/mp_router.py
View File

@@ -17,6 +17,7 @@ from services.mp_service import (
)
from utils import handle_minio_upload
from error_handlers import handle_general_error
from utils import settings
router = APIRouter(prefix="/mp", tags=["Material Project"])
logger = logging.getLogger(__name__)
@@ -55,10 +56,10 @@ async def search_from_material_project(request: Request):
os.remove(file_name)
# 格式化返回结果
res_chunk = "```json\n" + json.dumps(res[:5], indent=2) + "\n```"
res_chunk = "```json\n" + json.dumps(res[:settings.mp_topk], indent=2) + "\n```"
res_template = f"""
好的，以下是用户的查询结果：
由于返回长度的限制，我们只能返回前5个结果。如下：
由于返回长度的限制，我们只能返回前{settings.mp_topk}个结果。如下：
{res_chunk}
如果用户需要更多的结果，请提示用户修改查询条件，或者尝试使用其他查询参数。
同时我们将全部的的查询结果上传到MinIO中，请你提示用户可以通过以下链接下载：

1
router/oqmd_router.py
View File

@@ -53,5 +53,6 @@ def format_response(basic_data: list, table_data: str, phase_data: str) -> str:
response += "\n### Phase Diagram\n\n"
response += f"![Phase Diagram]({phase_data})\n\n"
response += "\n### Compounds at this composition\n\n"
response += "I highly recommend that you pass the phase map URL to the user wrapped in Markdown's image syntax!"
response += f"{table_data}\n"
return response

16
services/fairchem_service.py
View File

@@ -64,8 +64,20 @@ def optimize_structure(atoms: Atoms, output_format: str):
atoms.calc = calc
try:
import io
from contextlib import redirect_stdout
# 创建StringIO对象捕获输出
f = io.StringIO()
dyn = FIRE(FrechetCellFilter(atoms))
dyn.run(fmax=settings.fmax)
# 同时捕获并输出到控制台
with redirect_stdout(f):
dyn.run(fmax=settings.fmax)
# 获取捕获的日志
optimization_log = f.getvalue()
# 同时输出到控制台
print(optimization_log)
total_energy = atoms.get_total_energy()
# 处理对称性
@@ -87,6 +99,6 @@ def optimize_structure(atoms: Atoms, output_format: str):
# 上传到MinIO
url = handle_minio_upload(tmp_path, file_name)
return total_energy, content, url
return total_energy, content, optimization_log, url
finally:
os.unlink(tmp_path)

233
test_client.py Normal file
View File

@@ -0,0 +1,233 @@
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"))

1
utils.py
View File

@@ -17,6 +17,7 @@ class Settings(BaseSettings):
# Material Project
mp_api_key: Optional[str] = Field(None, env="MP_API_KEY")
mp_endpoint: Optional[str] = Field(None, env="MP_ENDPOINT")
mp_topk: Optional[int] = Field(3, env="MP_TOPK")
# Proxy
http_proxy: Optional[str] = Field(None, env="HTTP_PROXY")