HuiBrain/app/config.py

"""
配置管理模块
使用 pydantic-settings 从环境变量 / .env 文件加载配置
"""

from __future__ import annotations

from enum import Enum
from pathlib import Path
from typing import Optional

from pydantic import Field
from pydantic_settings import BaseSettings, SettingsConfigDict


class EmbeddingProvider(str, Enum):
    """嵌入模型提供商"""
    OPENAI = "openai"
    ZHIPU = "zhipu"
    LOCAL_BGE = "local_bge"
    DASHSCOPE = "dashscope"
    MINIMAX = "minimax"


class OCRProvider(str, Enum):
    """OCR 识别提供商"""
    PADDLEOCR = "paddleocr"
    ALIYUN = "aliyun"
    TENCENT = "tencent"
    AUTO = "auto"
    DEEPSEEK = "deepseek"


class Settings(BaseSettings):
    """
    全局配置类，所有配置项均可通过环境变量覆盖。
    环境变量优先级：系统环境变量 > .env 文件 > 默认值
    """

    model_config = SettingsConfigDict(
        env_file=".env",
        env_file_encoding="utf-8",
        case_sensitive=False,
        extra="ignore",
    )

    # ──────────────────────────── 数据库 ────────────────────────────
    DATABASE_URL: str = Field(
        default="postgresql+asyncpg://postgres:postgres@db:5432/edu_brain",
        description="PostgreSQL 异步连接字符串",
    )

    # ──────────────────────────── 嵌入模型 ────────────────────────────
    EMBEDDING_PROVIDER: EmbeddingProvider = Field(
        default=EmbeddingProvider.LOCAL_BGE,
        description="嵌入模型提供商: openai / zhipu / local_bge / dashscope",
    )
    EMBEDDING_MODEL: str = Field(
        default="BAAI/bge-large-zh-v1.5",
        description="嵌入模型名称",
    )
    EMBEDDING_DIMENSIONS: int = Field(
        default=1024,
        description="嵌入向量维度",
    )

    # ── OpenAI ──
    OPENAI_API_KEY: Optional[str] = Field(default=None, description="OpenAI API Key")
    OPENAI_BASE_URL: Optional[str] = Field(
        default=None,
        description="OpenAI 兼容接口地址（可用于代理或第三方兼容服务）",
    )

    # ── 智谱 AI ──
    ZHIPU_API_KEY: Optional[str] = Field(default=None, description="智谱 AI API Key")

    # ── 阿里云 DashScope ──
    DASHSCOPE_API_KEY: Optional[str] = Field(
        default=None, description="阿里云 DashScope API Key"
    )

    # ── 本地 BGE 模型 ──
    LOCAL_BGE_MODEL_PATH: Optional[str] = Field(
        default=None,
        description="本地 BGE 模型路径，为空则自动从 HuggingFace 下载",
    )

    # ──────────────────────────── OCR ────────────────────────────
    OCR_PROVIDER: OCRProvider = Field(
        default=OCRProvider.AUTO,
        description="OCR 提供商: paddleocr / aliyun / tencent / auto",
    )

    # ── 阿里云 OCR ──
    ALIYUN_OCR_ACCESS_KEY: Optional[str] = Field(
        default=None, description="阿里云 OCR AccessKey ID"
    )
    ALIYUN_OCR_SECRET: Optional[str] = Field(
        default=None, description="阿里云 OCR AccessKey Secret"
    )

    # ── 腾讯云 OCR ──
    TENCENT_OCR_SECRET_ID: Optional[str] = Field(
        default=None, description="腾讯云 OCR SecretId"
    )
    TENCENT_OCR_SECRET_KEY: Optional[str] = Field(
        default=None, description="腾讯云 OCR SecretKey"
    )

    # ── MiniMax ──
    MINIMAX_API_KEY: Optional[str] = Field(default=None, description="MiniMax API Key")
    MINIMAX_BASE_URL: str = Field(
        default="https://api.minimaxi.com/v1",
        description="MiniMax API 地址",
    )
    MINIMAX_EMBEDDING_MODEL: str = Field(
        default="MiniMax-Text-Embedding",
        description="MiniMax 嵌入模型名称",
    )
    MINIMAX_CHAT_MODEL: str = Field(
        default="MiniMax-M2",
        description="MiniMax Chat 模型名称",
    )

    # ── DeepSeek ──
    DEEPSEEK_API_KEY: Optional[str] = Field(default=None, description="DeepSeek API Key")
    DEEPSEEK_BASE_URL: str = Field(
        default="https://api.deepseek.com",
        description="DeepSeek API 地址（自部署时改为你的服务器地址）",
    )
    DEEPSEEK_OCR_MODEL: str = Field(
        default="deepseek-chat",
        description="DeepSeek OCR 使用的模型名称",
    )

    # ── LLM（用于查询扩展、实体检测等） ──
    LLM_PROVIDER: str = Field(
        default="minimax",
        description="LLM 提供商: minimax / deepseek / openai",
    )

    # ──────────────────────────── 数据目录 ────────────────────────────
    DATA_DIR: str = Field(
        default="/app/data",
        description="数据根目录（ transcripts / images / exports 的父目录）",
    )

    # ──────────────────────────── 文本分块 ────────────────────────────
    CHUNK_SIZE: int = Field(default=300, description="文本分块大小（字符数）")
    CHUNK_OVERLAP: int = Field(default=50, description="分块重叠字符数")

    # ──────────────────────────── 图片搜索 ────────────────────────────
    SEARCH_LIMIT: int = Field(default=3, description="图片搜索默认返回数量(1-10)")
    JUDGE_BATCH_SIZE: int = Field(default=10, description="每次LLM判断的文章数量(1-50)")

    # ──────────────────────────── 服务端口 ────────────────────────────
    HOST: str = Field(default="0.0.0.0", description="服务监听地址")
    PORT: int = Field(default=8000, description="服务监听端口")
    DEBUG: bool = Field(default=False, description="调试模式")

    # ──────────────────────────── CORS ────────────────────────────
    CORS_ORIGINS: list[str] = Field(
        default=["*"],
        description="允许的跨域来源列表",
    )

    # ──────────────────────────── 企业微信智能机器人（WebSocket 长连接） ────────────────────────────
    WEWORK_BOT_ENABLED: bool = Field(
        default=False,
        description="是否启用企业微信智能机器人（独立进程运行）",
    )
    WEWORK_BOT_ID: Optional[str] = Field(
        default=None,
        description="企业微信智能机器人 ID",
    )
    WEWORK_BOT_SECRET: Optional[str] = Field(
        default=None,
        description="企业微信智能机器人 Secret",
    )
    WEWORK_BOT_SEARCH_LIMIT: int = Field(
        default=3,
        ge=1,
        le=10,
        description="企业微信机器人搜索返回图片数量上限",
    )
    EDUBRAIN_BASE_URL: str = Field(
        default="http://localhost:8765",
        description="EduBrain 服务地址（供企业微信机器人调用搜索接口）",
    )

    # ──────────────────────────── 属性方法 ────────────────────────────

    @property
    def transcripts_dir(self) -> Path:
        """直播文字稿目录"""
        return Path(self.DATA_DIR) / "transcripts"

    @property
    def images_dir(self) -> Path:
        """聊天截图目录"""
        return Path(self.DATA_DIR) / "images"

    @property
    def exports_dir(self) -> Path:
        """导出文件目录"""
        return Path(self.DATA_DIR) / "exports"

    def ensure_dirs(self) -> None:
        """确保所有数据目录存在"""
        for d in (self.transcripts_dir, self.images_dir, self.exports_dir):
            d.mkdir(parents=True, exist_ok=True)


# ──────────────────────────── 全局单例 ────────────────────────────

settings = Settings()