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

router/fairchem_router.py

@@ -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,

router/mp_router.py

@@ -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中，请你提示用户可以通过以下链接下载：

router/oqmd_router.py

@@ -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

services/fairchem_service.py

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

test_client.py

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

utils.py

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