重构代码

router/fairchem_router.py

@@ -0,0 +1,60 @@
+"""
+Author: Yutang LI
+Institution: SIAT-MIC
+Contact: yt.li2@siat.ac.cn
+"""
+
+from fastapi import APIRouter, Body, Query
+from fastapi.responses import JSONResponse
+import logging
+from error_handlers import handle_general_error
+from services.fairchem_service import (
+    init_model,
+    convert_structure,
+    optimize_structure
+)
+
+router = APIRouter(prefix="/fairchem", tags=["fairchem"])
+logger = logging.getLogger(__name__)
+
+# 初始化模型
+init_model()
+
+@router.post("/optimize_structure")
+async def optimize_structure_endpoint(
+    content: str = Body(..., description="Input structure content"),
+    input_format: str = Query("cif", description="Input format (cif, poscar, json, xyz)"),
+    output_format: str = Query("cif", description="Output format (cif, poscar, json, xyz)")
+):
+    try:
+        # 转换输入结构
+        atoms = convert_structure(input_format, content)
+        if atoms is None:
+            return JSONResponse(
+                status_code=400,
+                content={"status": "error", "data": f"Invalid {input_format} content"}
+            )
+        
+        # 优化结构
+        total_energy, optimized_content, 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
+
+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
+"""
+        return JSONResponse(
+            status_code=200,
+            content={"status": "success", "data": format_result}
+        )
+    
+    except Exception as e:
+        return handle_general_error(e)
+
+if __name__ == "__main__":
+    init_model()