config_reader.py 1.05 KB
Newer Older
kernel.h@qq.com's avatar
kernel.h@qq.com committed
1
2
3
4
"""
根据bucket的名字返回对应的s3 AK, SK,endpoint三元组

"""
5
6
7
8
9
import json
import os

from loguru import logger

kernel.h@qq.com's avatar
kernel.h@qq.com committed
10
11
12

def get_s3_config(bucket_name: str):
    """
赵小蒙's avatar
赵小蒙 committed
13
    ~/magic-pdf.json 读出来
kernel.h@qq.com's avatar
kernel.h@qq.com committed
14
    """
15

赵小蒙's avatar
赵小蒙 committed
16
17
18
    home_dir = os.path.expanduser("~")

    config_file = os.path.join(home_dir, "magic-pdf.json")
19
20

    if not os.path.exists(config_file):
赵小蒙's avatar
赵小蒙 committed
21
        raise Exception("magic-pdf.json not found")
22
23
24
25

    with open(config_file, "r") as f:
        config = json.load(f)

赵小蒙's avatar
赵小蒙 committed
26
27
28
    bucket_info = config.get("bucket_info")
    if bucket_name not in bucket_info:
        raise Exception("bucket_name not found in magic-pdf.json")
29

赵小蒙's avatar
赵小蒙 committed
30
    access_key, secret_key, storage_endpoint = bucket_info[bucket_name]
31

赵小蒙's avatar
赵小蒙 committed
32
    if access_key is None or secret_key is None or storage_endpoint is None:
赵小蒙's avatar
赵小蒙 committed
33
        raise Exception("ak, sk or endpoint not found in magic-pdf.json")
34

赵小蒙's avatar
赵小蒙 committed
35
    # logger.info(f"get_s3_config: ak={access_key}, sk={secret_key}, endpoint={storage_endpoint}")
36

赵小蒙's avatar
赵小蒙 committed
37
    return access_key, secret_key, storage_endpoint
38
39
40
41


if __name__ == '__main__':
    ak, sk, endpoint = get_s3_config("llm-raw")