config.py 3.5 KB
"""
Query Engine 配置管理模块

此模块使用 pydantic-settings 管理 Query Engine 的配置,支持从环境变量和 .env 文件自动加载。
数据模型定义位置:
- 本文件 - 配置模型定义
"""

from pathlib import Path
from pydantic_settings import BaseSettings
from pydantic import Field
from typing import Optional
from loguru import logger


# 计算 .env 优先级:优先当前工作目录,其次项目根目录
PROJECT_ROOT: Path = Path(__file__).resolve().parents[2]
CWD_ENV: Path = Path.cwd() / ".env"
ENV_FILE: str = str(CWD_ENV if CWD_ENV.exists() else (PROJECT_ROOT / ".env"))


class Settings(BaseSettings):
    """
    Query Engine 全局配置;支持 .env 和环境变量自动加载。
    变量名与原 config.py 大写一致,便于平滑过渡。
    """
    
    # ======================= LLM 相关 =======================
    QUERY_ENGINE_API_KEY: str = Field(..., description="Query Engine LLM API密钥,用于主LLM。您可以更改每个部分LLM使用的API,🚩只要兼容OpenAI请求格式都可以,定义好KEY、BASE_URL与MODEL_NAME即可正常使用。")
    QUERY_ENGINE_BASE_URL: Optional[str] = Field(None, description="Query Engine LLM接口BaseUrl,可自定义厂商API")
    QUERY_ENGINE_MODEL_NAME: str = Field(..., description="Query Engine LLM模型名称")
    QUERY_ENGINE_PROVIDER: Optional[str] = Field(None, description="Query Engine LLM提供商(兼容字段)")
    
    # ================== 网络工具配置 ====================
    TAVILY_API_KEY: str = Field(..., description="Tavily API(申请地址:https://www.tavily.com/)API密钥,用于Tavily网络搜索")
    
    # ================== 搜索参数配置 ====================
    SEARCH_TIMEOUT: int = Field(240, description="搜索超时(秒)")
    SEARCH_CONTENT_MAX_LENGTH: int = Field(20000, description="用于提示的最长内容长度")
    MAX_REFLECTIONS: int = Field(2, description="最大反思轮数")
    MAX_PARAGRAPHS: int = Field(5, description="最大段落数")
    MAX_SEARCH_RESULTS: int = Field(20, description="最大搜索结果数")
    
    # ================== 输出配置 ====================
    OUTPUT_DIR: str = Field("reports", description="输出目录")
    SAVE_INTERMEDIATE_STATES: bool = Field(True, description="是否保存中间状态")
    
    class Config:
        env_file = ENV_FILE
        env_prefix = ""
        case_sensitive = False
        extra = "allow"


# 创建全局配置实例
settings = Settings()

def print_config(config: Settings):
    """
    打印配置信息
    
    Args:
        config: Settings配置对象
    """
    message = ""
    message += "=== Query Engine 配置 ===\n"
    message += f"LLM 模型: {config.QUERY_ENGINE_MODEL_NAME}\n"
    message += f"LLM Base URL: {config.QUERY_ENGINE_BASE_URL or '(默认)'}\n"
    message += f"Tavily API Key: {'已配置' if config.TAVILY_API_KEY else '未配置'}\n"
    message += f"搜索超时: {config.SEARCH_TIMEOUT} 秒\n"
    message += f"最长内容长度: {config.SEARCH_CONTENT_MAX_LENGTH}\n"
    message += f"最大反思次数: {config.MAX_REFLECTIONS}\n"
    message += f"最大段落数: {config.MAX_PARAGRAPHS}\n"
    message += f"最大搜索结果数: {config.MAX_SEARCH_RESULTS}\n"
    message += f"输出目录: {config.OUTPUT_DIR}\n"
    message += f"保存中间状态: {config.SAVE_INTERMEDIATE_STATES}\n"
    message += f"LLM API Key: {'已配置' if config.QUERY_ENGINE_API_KEY else '未配置'}\n"
    message += "========================\n"
    logger.info(message)