104 lines
3.3 KiB
Python
104 lines
3.3 KiB
Python
"""
|
|
Author: Yutang LI
|
|
Institution: SIAT-MIC
|
|
Contact: yt.li2@siat.ac.cn
|
|
"""
|
|
|
|
import os
|
|
import boto3
|
|
import logging
|
|
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")
|
|
mp_topk: Optional[int] = Field(3, env="MP_TOPK")
|
|
|
|
# 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()
|