config_reader.py 1.32 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):
    """
13
    ~/magic_pdf_config.json 读出来
kernel.h@qq.com's avatar
kernel.h@qq.com committed
14
    """
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
    if os.name == "posix":  # Linux or macOS
        home_dir = os.path.expanduser("~")
    elif os.name == "nt":  # Windows
        home_dir = os.path.expandvars("%USERPROFILE%")
    else:
        raise Exception("Unsupported operating system")

    config_file = os.path.join(home_dir, "magic_pdf_config.json")

    if not os.path.exists(config_file):
        raise Exception("magic_pdf_config.json not found")

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

    if bucket_name not in config:
        raise Exception("bucket_name not found in magic_pdf_config.json")

赵小蒙's avatar
赵小蒙 committed
33
34
35
    access_key = config[bucket_name].get("ak")
    secret_key = config[bucket_name].get("sk")
    storage_endpoint = config[bucket_name].get("endpoint")
36

赵小蒙's avatar
赵小蒙 committed
37
    if access_key is None or secret_key is None or storage_endpoint is None:
38
39
        raise Exception("ak, sk or endpoint not found in magic_pdf_config.json")

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

赵小蒙's avatar
赵小蒙 committed
42
    return access_key, secret_key, storage_endpoint
43
44
45
46


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