重构代码

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

router/mp_router.py

@@ -0,0 +1,70 @@
+"""
+Author: Yutang LI
+Institution: SIAT-MIC
+Contact: yt.li2@siat.ac.cn
+"""
+
+import os
+from fastapi import APIRouter, Request
+import json
+import logging
+import datetime
+from typing import Dict
+from services.mp_service import (
+    parse_search_parameters,
+    process_search_results,
+    execute_search
+)
+from utils import handle_minio_upload
+from error_handlers import handle_general_error
+
+router = APIRouter(prefix="/mp", tags=["Material Project"])
+logger = logging.getLogger(__name__)
+
+@router.get("/search")
+async def search_from_material_project(request: Request):
+    # 打印请求日志
+    logger.info(f"Received request: {request.method} {request.url}")
+    logger.info(f"Query parameters: {request.query_params}")
+
+    try:
+        # 解析查询参数
+        search_args = parse_search_parameters(request.query_params)
+        
+        # 执行搜索
+        docs = await execute_search(search_args)
+        
+        # 处理搜索结果
+        res = process_search_results(docs)
+        
+        if len(res) == 0:
+            return {"status": "success", "data": "No results found, please try again."}
+        
+        # 上传结果到MinIO
+        timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+        file_name = f"mp_search_results_{timestamp}.json"
+        
+        # 将结果写入临时文件
+        with open(file_name, 'w') as f:
+            json.dump(res, f, indent=2)
+        
+        # 上传并获取URL
+        url = handle_minio_upload(file_name, file_name)
+        
+        # 删除临时文件
+        os.remove(file_name)
+        
+        # 格式化返回结果
+        res_chunk = "```json\n" + json.dumps(res[:5], indent=2) + "\n```"
+        res_template = f"""
+好的，以下是用户的查询结果：
+由于返回长度的限制，我们只能返回前5个结果。如下：
+{res_chunk}
+如果用户需要更多的结果，请提示用户修改查询条件，或者尝试使用其他查询参数。
+同时我们将全部的的查询结果上传到MinIO中，请你提示用户可以通过以下链接下载：
+[Download]({url})
+"""
+        return {"status": "success", "data": res_template}
+        
+    except Exception as e:
+        return handle_general_error(e)

router/oqmd_router.py

@@ -0,0 +1,57 @@
+"""
+Author: Yutang LI
+Institution: SIAT-MIC
+Contact: yt.li2@siat.ac.cn
+"""
+import os
+from fastapi import APIRouter, Request
+from fastapi.responses import JSONResponse
+import logging
+from error_handlers import handle_general_error
+from services.oqmd_service import (
+    fetch_oqmd_data,
+    parse_oqmd_html,
+    render_and_save_charts
+)
+
+router = APIRouter(prefix="/oqmd", tags=["OQMD"])
+logger = logging.getLogger(__name__)
+
+@router.get("/search")
+async def search_from_oqmd_by_composition(request: Request):
+    """通过成分搜索OQMD数据"""
+    try:
+        # 打印请求日志
+        logger.info(f"Received request: {request.method} {request.url}")
+        logger.info(f"Query parameters: {request.query_params}")
+
+        # 获取并解析数据
+        composition = request.query_params['composition']
+        html = await fetch_oqmd_data(composition)
+        basic_data, table_data, phase_data = parse_oqmd_html(html)
+        
+        # 渲染并保存图表
+        phase_diagram_url = await render_and_save_charts(phase_data)
+        
+        # 返回格式化后的响应
+        return JSONResponse(
+            status_code=200,
+            content={
+                "status": "success",
+                "data": format_response(basic_data, table_data, phase_diagram_url)
+            }
+        )
+    
+    except Exception as e:
+        return handle_general_error(e)
+
+def format_response(basic_data: list, table_data: str, phase_data: str) -> str:
+    """格式化响应数据"""
+    response = "### OQMD Data\n"
+    for item in basic_data:
+        response += f"**{item}**\n"
+    response += "\n### Phase Diagram\n\n"
+    response += f"![Phase Diagram]({phase_data})\n\n"
+    response += "\n### Compounds at this composition\n\n"
+    response += f"{table_data}\n"
+    return response