ailine/backend/app/config.py

"""
环境变量集中管理模块
所有配置项统一定义，避免散落在各个文件中
配置分组：相关配置放在一起，URL 和 API Key 配对
所有配置直接从环境变量读取，无默认值，避免配置混乱
需要类型转换的配置在此处理
"""

import os
from dotenv import load_dotenv

load_dotenv()

# ========== 辅助函数：类型转换 ==========
def _get_str(key: str) -> str | None:
    """获取字符串配置"""
    return os.getenv(key)


def _get_int(key: str) -> int | None:
    """获取整数配置，自动转换"""
    value = os.getenv(key)
    if value is not None:
        try:
            return int(value)
        except (ValueError, TypeError):
            pass
    return None


def _get_bool(key: str) -> bool | None:
    """获取布尔配置，自动转换"""
    value = os.getenv(key)
    if value is not None:
        return value.lower() in ("true", "1", "yes", "on")
    return None


# ========== 第三方 API 密钥 ==========
ZHIPUAI_API_KEY = _get_str("ZHIPUAI_API_KEY")
DEEPSEEK_API_KEY = _get_str("DEEPSEEK_API_KEY")
SILICONFLOW_API_KEY = _get_str("SILICONFLOW_API_KEY")
BAOSI_API_KEY = _get_str("BAOSI_API_KEY")


# ========== 智谱 API 配置 ==========
# 嵌入模型：根据 https://docs.bigmodel.cn/cn/guide/start/model-overview
# 可选：embedding-2 (1024维)、embedding-3 (2048维)
# 注意：如果 Qdrant collection 是1024维，请使用 embedding-2
ZHIPU_EMBEDDING_MODEL = _get_str("ZHIPU_EMBEDDING_MODEL") or "embedding-2"
# 重排模型：可选 rerank-1、rerank-2
ZHIPU_RERANK_MODEL = _get_str("ZHIPU_RERANK_MODEL") or "rerank-2"
ZHIPU_API_BASE = _get_str("ZHIPU_API_BASE") or "https://open.bigmodel.cn/api/paas/v4"


# ========== 硅基流动(SiliconFlow) API 配置 ==========
# 重排模型：BAAI/bge-reranker-v2-m3
SILICONFLOW_RERANK_MODEL = _get_str("SILICONFLOW_RERANK_MODEL") or "BAAI/bge-reranker-v2-m3"
SILICONFLOW_API_BASE = _get_str("SILICONFLOW_API_BASE") or "https://api.siliconflow.cn/v1"


# ========== Baosi API 配置 ==========
BAOSI_API_BASE = _get_str("BAOSI_API_BASE") or "https://api.baosiapi.com"
BAOSI_MODEL = _get_str("BAOSI_MODEL") or "ops4.7"


# ========== 稀疏模型配置 ==========
SPARSE_MODEL_PATH = _get_str("SPARSE_MODEL_PATH") or "./models/sparse"
SPARSE_MODEL_NAME = _get_str("SPARSE_MODEL_NAME") or "Qdrant/bm25"
FASTEMBED_CACHE_PATH = _get_str("FASTEMBED_CACHE_PATH") or "./models/fastembed_cache"

# ========== 本地模型配置 ==========
LOCAL_MODEL_NAME = _get_str("LOCAL_MODEL_NAME") or "gemma-4-E4B-it"

# ========== llama.cpp 服务配置（URL + API密钥 配对） ==========
# 主 LLM 服务
VLLM_BASE_URL = _get_str("VLLM_BASE_URL")
LLM_API_KEY = _get_str("LLM_API_KEY")

# Embedding 服务 (用于 Mem0 的向量化)
LLAMACPP_EMBEDDING_URL = _get_str("LLAMACPP_EMBEDDING_URL")
LLAMACPP_API_KEY = _get_str("LLAMACPP_API_KEY")

# Reranker 服务
LLAMACPP_RERANKER_URL = _get_str("LLAMACPP_RERANKER_URL")

# ========== 小模型配置（查询改写、意图分类等简单任务） ==========
# 默认复用大模型配置，后续可单独配置
# 本地小模型（默认复用 VLLM 配置）
SMALL_VLLM_BASE_URL = _get_str("SMALL_VLLM_BASE_URL")
SMALL_LLM_API_KEY = _get_str("SMALL_LLM_API_KEY")
SMALL_LOCAL_MODEL_NAME = _get_str("SMALL_LOCAL_MODEL_NAME") or LOCAL_MODEL_NAME
# 如果小模型没单独配置，用大模型的配置
if not SMALL_VLLM_BASE_URL:
    SMALL_VLLM_BASE_URL = VLLM_BASE_URL
if not SMALL_LLM_API_KEY:
    SMALL_LLM_API_KEY = LLM_API_KEY

# DeepSeek 小模型（默认复用 DeepSeek 配置）
SMALL_DEEPSEEK_API_KEY = _get_str("SMALL_DEEPSEEK_API_KEY")
SMALL_DEEPSEEK_MODEL = _get_str("SMALL_DEEPSEEK_MODEL") or "deepseek-chat"
SMALL_DEEPSEEK_API_BASE = _get_str("SMALL_DEEPSEEK_API_BASE") or "https://api.deepseek.com"
# 如果小模型没单独配置，用大模型的配置
if not SMALL_DEEPSEEK_API_KEY:
    SMALL_DEEPSEEK_API_KEY = DEEPSEEK_API_KEY


# ========== Qdrant 向量数据库配置（URL + API密钥 配对） ==========
QDRANT_URL = _get_str("QDRANT_URL")
QDRANT_API_KEY = _get_str("QDRANT_API_KEY")
QDRANT_COLLECTION_NAME = _get_str("QDRANT_COLLECTION_NAME")


# ========== PostgreSQL 数据库配置（分离配置 + 完整URI） ==========
# 分离配置（优先使用）
DB_HOST = _get_str("DB_HOST")
DB_PORT = _get_int("DB_PORT")
DB_USER = _get_str("DB_USER")
DB_PASSWORD = _get_str("DB_PASSWORD")
DB_NAME = _get_str("DB_NAME")

# 完整连接字符串（直接从环境变量读取）
DB_URI = _get_str("DB_URI")
if not DB_URI and all([DB_HOST, DB_PORT, DB_USER, DB_PASSWORD, DB_NAME]):
    DB_URI = f"postgresql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}?sslmode=disable"

# ========== 后端服务配置 ==========
BACKEND_PORT = _get_int("BACKEND_PORT")


# ========== Mem0 记忆层配置 ==========
# 记忆提取间隔：每 N 轮对话生成一次摘要
MEMORY_SUMMARIZE_INTERVAL = _get_int("MEMORY_SUMMARIZE_INTERVAL")


# ========== Tavily 搜索配置 ==========
# Tavily API：https://app.tavily.com
# 免费额度：1000次/天
TAVILY_API_KEY = _get_str("TAVILY_API_KEY")
TAVILY_MAX_RESULTS = _get_int("TAVILY_MAX_RESULTS") or 5


# ========== Graph 执行追踪配置 ==========
# 是否启用 Graph 流转追踪（通过环境变量控制）
ENABLE_GRAPH_TRACE = _get_bool("ENABLE_GRAPH_TRACE")


# ========== 日志配置 ==========
LOG_LEVEL = _get_str("LOG_LEVEL")
DEBUG = _get_bool("DEBUG")