重构代码

utils.py

@@ -1,8 +1,102 @@
+"""
+Author: Yutang LI
+Institution: SIAT-MIC
+Contact: yt.li2@siat.ac.cn
+"""
+
+import os
+import boto3
 import logging
-from multiprocessing import Process, Manager
-import asyncio
-from typing import Dict, Any, List
-from mp_api.client import MPRester
+from typing import Optional
+from pydantic import Field
+from pydantic_settings import BaseSettings
 
 logger = logging.getLogger(__name__)
 
+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")
+    
+    # Proxy
+    http_proxy: Optional[str] = Field(None, env="HTTP_PROXY")
+    https_proxy: Optional[str] = Field(None, env="HTTPS_PROXY")
+    
+    # FairChem
+    fairchem_model_path: Optional[str] = Field(None, env="FAIRCHEM_MODEL_PATH")
+    fmax: Optional[float] = Field(0.05, env="FMAX")
+    
+    # MinIO
+    minio_endpoint: Optional[str] = Field(None, env="MINIO_ENDPOINT")
+    internal_minio_endpoint: Optional[str] = Field(None, env="INTERNAL_MINIO_ENDPOINT")
+    minio_access_key: Optional[str] = Field(None, env="MINIO_ACCESS_KEY")
+    minio_secret_key: Optional[str] = Field(None, env="MINIO_SECRET_KEY")
+    minio_bucket: Optional[str] = Field("mars-toolkit", env="MINIO_BUCKET")
+    
+    class Config:
+        env_file = ".env"
+        env_file_encoding = "utf-8"
+
+def get_minio_client(settings: Settings):
+    """获取MinIO客户端"""
+    return boto3.client(
+        's3',
+        endpoint_url=settings.internal_minio_endpoint or settings.minio_endpoint,
+        aws_access_key_id=settings.minio_access_key,
+        aws_secret_access_key=settings.minio_secret_key
+    )
+
+def handle_minio_upload(file_path: str, file_name: str) -> str:
+    """统一处理MinIO上传"""
+    try:
+        client = get_minio_client(settings)
+        client.upload_file(file_path, settings.minio_bucket, file_name, ExtraArgs={"ACL": "private"})
+        
+        # 生成预签名 URL
+        url = client.generate_presigned_url(
+            'get_object',
+            Params={'Bucket': settings.minio_bucket, 'Key': file_name},
+            ExpiresIn=3600
+        )
+        return url.replace(settings.internal_minio_endpoint or "", settings.minio_endpoint)
+    except Exception as e:
+        from error_handlers import handle_minio_error
+        return handle_minio_error(e)
+
+def setup_logging():
+    """配置日志记录"""
+    logging.config.dictConfig({
+        'version': 1,
+        'disable_existing_loggers': False,
+        'formatters': {
+            'standard': {
+                'format': '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+                'datefmt': '%Y-%m-%d %H:%M:%S'
+            },
+        },
+        'handlers': {
+            'console': {
+                'level': 'INFO',
+                'class': 'logging.StreamHandler',
+                'formatter': 'standard'
+            },
+            'file': {
+                'level': 'DEBUG',
+                'class': 'logging.handlers.RotatingFileHandler',
+                'filename': 'mars_toolkit.log',
+                'maxBytes': 10485760,  # 10MB
+                'backupCount': 5,
+                'formatter': 'standard'
+            }
+        },
+        'loggers': {
+            '': {
+                'handlers': ['console', 'file'],
+                'level': 'INFO',
+                'propagate': True
+            }
+        }
+    })
+
+# 初始化配置
+settings = Settings()