jude-cehua/config.py

import os
from dataclasses import dataclass
from typing import Optional, Tuple
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

@dataclass
class AppConfig:
    """应用配置数据类"""
    # API配置
    api_key: str
    app_id: str
    api_endpoint: str
    
    # 应用配置
    secret_key: str
    port: int
    debug: bool
    
    # 数据库配置
    db_host: str
    db_port: int
    db_name: str
    db_user: str
    db_password: str
    db_charset: str
    
    # 安全配置
    max_message_length: int
    max_conversation_history: int
    rate_limit_requests: int
    rate_limit_window: int
    
    # 日志配置
    log_level: str
    log_format: str
    
    @classmethod
    def from_env(cls) -> 'AppConfig':
        """从环境变量创建配置"""
        return cls(
            # API配置
            api_key=os.getenv("API_KEY") or "",
            app_id=os.getenv("APP_ID") or "",
            api_endpoint=os.getenv("API_ENDPOINT") or "",
            
            # 应用配置
            secret_key=os.getenv("SECRET_KEY") or os.urandom(24).hex(),
            port=int(os.getenv("PORT", "8000")),
            debug=os.getenv("FLASK_DEBUG", "false").lower() == "true",
            
            # 数据库配置
            db_host=os.getenv("DB_HOST", "localhost"),
            db_port=int(os.getenv("DB_PORT", "3306")),
            db_name=os.getenv("DB_NAME", "ai_writing_db"),
            db_user=os.getenv("DB_USER", "root"),
            db_password=os.getenv("DB_PASSWORD", ""),
            db_charset=os.getenv("DB_CHARSET", "utf8mb4"),
            
            # 安全配置
            max_message_length=int(os.getenv("MAX_MESSAGE_LENGTH", "20000")),
            max_conversation_history=int(os.getenv("MAX_CONVERSATION_HISTORY", "50")),
            rate_limit_requests=int(os.getenv("RATE_LIMIT_REQUESTS", "100")),
            rate_limit_window=int(os.getenv("RATE_LIMIT_WINDOW", "3600")),
            
            # 日志配置
            log_level=os.getenv("LOG_LEVEL", "INFO"),
            log_format=os.getenv("LOG_FORMAT", "%(asctime)s - %(name)s - %(levelname)s - %(message)s")
        )
    
    def validate(self) -> None:
        """验证配置的有效性"""
        if not self.api_key:
            raise ValueError("API_KEY 环境变量必须设置")
        # 允许通过基础 APP_ID 或 按智能体区分的 APP_ID_{AGENT} 提供
        has_base_app_id = bool(self.app_id)
        has_any_agent_app_id = any(k.startswith("APP_ID_") and os.getenv(k) for k in os.environ.keys())
        if not (has_base_app_id or has_any_agent_app_id):
            raise ValueError("必须设置 APP_ID，或通过 APP_ID_{AGENT} 之一提供，例如 APP_ID_KUAISHOU")
        if not self.api_endpoint:
            raise ValueError("API_ENDPOINT 环境变量必须设置")
        
        if self.max_message_length <= 0:
            raise ValueError("MAX_MESSAGE_LENGTH 必须大于0")
        if self.max_conversation_history <= 0:
            raise ValueError("MAX_CONVERSATION_HISTORY 必须大于0")
        if self.rate_limit_requests <= 0:
            raise ValueError("RATE_LIMIT_REQUESTS 必须大于0")
        if self.rate_limit_window <= 0:
            raise ValueError("RATE_LIMIT_WINDOW 必须大于0")

    # 新增：根据智能体名称获取 APP_ID（默认快手），带来源信息
    def get_app_id_for_agent_with_source(self, agent_name: Optional[str]) -> Tuple[str, str]:
        """
        根据前端传入的智能体标识，从环境变量读取对应的 APP_ID，并返回使用的环境变量键名。
        优先匹配：
        其次再尝试 APP_ID_{ALIAS} 以及其它别名；最后回退 APP_ID（通用）与 APP_ID_KUAISHOU。
        返回: (app_id, source_env_key)
        若找不到返回 ("", "")
        """
        normalized = (agent_name or "kuaishou").strip().lower()

        # 1) 按你的明确映射优先
        primary_map = {
            'kuaishou': ["APP_ID"],
            'toutiao': ["TOUTIAO_APP_ID"],
            'pathology': ["PATHOLOGY_APP_ID"], #病理分析
            'planning': ["PLANNING_APP_ID"], #策划大师
            'diagnosis': ["DIAGNOSIS_APP_ID"], #投放诊断
        }

        # 2) 兼容性别名（作为次优先）
        alias_map = {
            'kuaishou': ["APP_ID_KUAISHOU", "APP_ID_KS",],
            'toutiao': ["APP_ID_TOUTIAO", "APP_ID_TT", ],

            'planning': ["APP_ID_PLANNING", "APP_ID_PL",], #病理分析
            'pathology': ["APP_ID_PATHOLOGY",], #策划大师
            'diagnosis': ["DIAGNOSIS_APP_ID",], #投放诊断
        }

        for key in primary_map.get(normalized, []):
            val = os.getenv(key)
            if val:
                return val, key

        # 生成 APP_ID_{AGENT.upper()} 形式作为额外兼容
        generic_key = f"APP_ID_{normalized.upper()}"
        if generic_key not in alias_map.get(normalized, []):
            alias_candidates = [generic_key] + alias_map.get(normalized, [])
        else:
            alias_candidates = alias_map.get(normalized, [])

        for key in alias_candidates:
            val = os.getenv(key)
            if val:
                return val, key

        # 3) 通用 APP_ID 兜底（非快手也可作为兜底）
        if self.app_id:
            return self.app_id, "APP_ID"
        fallback = os.getenv("APP_ID_KUAISHOU")
        if fallback:
            return fallback, "APP_ID_KUAISHOU"
        return "", ""

    # 兼容旧方法：仅返回 APP_ID
    def get_app_id_for_agent(self, agent_name: Optional[str]) -> str:
        app_id, _ = self.get_app_id_for_agent_with_source(agent_name)
        return app_id
    
    def get_database_url(self) -> str:
        """获取数据库连接URL"""
        return f"mysql+pymysql://{self.db_user}:{self.db_password}@{self.db_host}:{self.db_port}/{self.db_name}?charset={self.db_charset}"

# 创建全局配置实例
config = AppConfig.from_env()

# 验证配置
try:
    config.validate()
except ValueError as e:
    exit(1)