self-constructing_graph/config/site.py

import os
from dotenv import load_dotenv
from urllib.parse import quote
from pydantic_settings import BaseSettings
from typing import Dict, Union, Optional

load_dotenv()


class SiteConfig:
    def __init__(self):
        self.load_config()
    
    def load_config(self):        
        self.config = {
            "SITE_NAME": os.getenv("SITE_NAME", "DEMO"),
            "SITE_DESCRIPTION": os.getenv("SITE_DESCRIPTION", "ChatGPT"),
            "SITE_URL": os.getenv("SITE_URL", ""),
            "SITE_LOGO": os.getenv("SITE_LOGO", ""),
            "SITE_FAVICON": os.getenv("SITE_FAVICON"),
            'ELASTICSEARCH_HOST': os.getenv("ELASTICSEARCH_HOST"),
            'ELASTICSEARCH_USER': os.getenv("ELASTICSEARCH_USER"),
            'ELASTICSEARCH_PWD': os.getenv("ELASTICSEARCH_PWD"),
            'WORD_INDEX': os.getenv("WORD_INDEX"),
            'TITLE_INDEX': os.getenv("TITLE_INDEX"),
            'CHUNC_INDEX': os.getenv("CHUNC_INDEX"),
            'DEEPSEEK_API_URL': os.getenv("DEEPSEEK_API_URL"),
            'DEEPSEEK_API_KEY': os.getenv("DEEPSEEK_API_KEY"),
            'CACHED_DATA_PATH': os.getenv("CACHED_DATA_PATH"),
            'UPDATE_DATA_PATH': os.getenv("UPDATE_DATA_PATH"),
            'FACTOR_DATA_PATH': os.getenv("FACTOR_DATA_PATH"),
            'GRAPH_API_URL': os.getenv("GRAPH_API_URL"),
            'EMBEDDING_MODEL': os.getenv("EMBEDDING_MODEL"),
            'DOC_PATH': os.getenv("DOC_PATH"),
            'DOC_STORAGE_PATH': os.getenv("DOC_STORAGE_PATH"),
            'TRUNC_OUTPUT_PATH': os.getenv("TRUNC_OUTPUT_PATH"),
            'DOC_ABSTRACT_OUTPUT_PATH': os.getenv("DOC_ABSTRACT_OUTPUT_PATH"),
            'JIEBA_USER_DICT': os.getenv("JIEBA_USER_DICT"),
            'JIEBA_STOP_DICT': os.getenv("JIEBA_STOP_DICT"),
            'POSTGRESQL_HOST':  os.getenv("POSTGRESQL_HOST","localhost"),
            'POSTGRESQL_DATABASE':  os.getenv("POSTGRESQL_DATABASE","kg"),
            'POSTGRESQL_USER':  os.getenv("POSTGRESQL_USER","dify"),
            'POSTGRESQL_PASSWORD':  os.getenv("POSTGRESQL_PASSWORD",quote("difyai123456")),
            'JOB_PATH': os.getenv("JOB_PATH","/app/agent/jobs"),
        }
    def get_config(self, config_name, default=None): 
        config_name = config_name.upper()     
        value = os.getenv(config_name, None)  
        if value:
            return value
        
        if config_name in self.config:            
            return self.config[config_name]
        else:
            return default
    def check_config(self, config_list):
        for item in config_list:
            if not self.get_config(item):
                raise ValueError(f"Configuration '{item}' is not set.")


class Settings(BaseSettings):
    # PostgreSQL配置
    DATABASE_URL: str = "postgresql://knowledge:qwer1234.@173.18.12.203:5432/knowledge_base"

    # MinIO配置
    MINIO_ENDPOINT: str = "173.18.12.199:9000"
    MINIO_ACCESS_KEY: str = "yvhNcezRwQvPuUylHqrg"
    MINIO_SECRET_KEY: str = "QQCejGeENpUIkGr4yfDaubPwWCoV29xHoXv6gHYU"
    MINIO_BUCKET_NAME: str = "knowledge-base"
    MINIO_SECURE: bool = False

    # 文件上传配置
    MAX_FILE_COUNT: int = 100
    ALLOWED_EXTENSIONS: Dict[str, Dict[str, Union[int, Optional[int]]]] = {
        "doc": {"max_size": 50, "max_pages": 1000},
        "txt": {"max_size": 10, "max_pages": None},
        "docx": {"max_size": 50, "max_pages": 1000},
        "pdf": {"max_size": 500, "max_pages": 3000},
        "ppt": {"max_size": 50, "max_pages": 1000},
        "pptx": {"max_size": 50, "max_pages": 1000},
        "md": {"max_size": 10, "max_pages": None}
    }


settings = Settings()