import json import os from pathlib import Path import yaml from src.utils.logging_config import logger class SimpleConfig(dict): def __key(self, key): return "" if key is None else key # 目前忘记了这里为什么要 lower 了，只能说配置项最好不要有大写的 def __str__(self): return json.dumps(self) def __setattr__(self, key, value): self[self.__key(key)] = value def __getattr__(self, key): return self.get(self.__key(key)) def __getitem__(self, key): return self.get(self.__key(key)) def __setitem__(self, key, value): return super().__setitem__(self.__key(key), value) def __dict__(self): return {k: v for k, v in self.items()} def update(self, other): for key, value in other.items(): self[key] = value class Config(SimpleConfig): def __init__(self): super().__init__() self._config_items = {} self.save_dir = os.getenv("SAVE_DIR", "saves") self.filename = str(Path(f"{self.save_dir}/config/base.yaml")) os.makedirs(os.path.dirname(self.filename), exist_ok=True) self._models_config_path: Path | None = os.getenv("OVERRIDE_DEFAULT_MODELS_CONFIG_WITH") self._update_models_from_file() ### >>> 默认配置 # 功能选项 self.add_item("enable_reranker", default=False, des="是否开启重排序") self.add_item("enable_content_guard", default=False, des="是否启用内容审查") self.add_item("enable_content_guard_llm", default=False, des="是否启用LLM内容审查") self.add_item( "content_guard_llm_model", default="siliconflow/Qwen/Qwen3-235B-A22B-Instruct-2507", des="内容审查LLM模型" ) # 默认智能体配置 self.add_item("default_agent_id", default="", des="默认智能体ID") # 模型配置 ## 注意这里是模型名，而不是具体的模型路径，默认使用 HuggingFace 的路径 ## 如果需要自定义本地模型路径，则在 .env 中配置 MODEL_DIR self.add_item("model_provider", default="siliconflow", des="模型提供商", choices=list(self.model_names.keys())) self.add_item("model_name", default="zai-org/GLM-4.5", des="模型名称") self.add_item( "fast_model", default="siliconflow/THUDM/GLM-4-9B-0414", des="快速响应模型", ) self.add_item( "embed_model", default="siliconflow/BAAI/bge-m3", des="Embedding 模型", choices=list(self.embed_model_names.keys()), ) self.add_item( "reranker", default="siliconflow/BAAI/bge-reranker-v2-m3", des="Re-Ranker 模型", choices=list(self.reranker_names.keys()), ) # noqa: E501 ### <<< 默认配置结束 self.load() self.handle_self() def add_item(self, key, default, des=None, choices=None): self.__setattr__(key, default) self._config_items[key] = {"default": default, "des": des, "choices": choices} def __dict__(self): blocklist = [ "_config_items", "model_names", "model_provider_status", "embed_model_names", "reranker_names", "_models_config_path", ] return {k: v for k, v in self.items() if k not in blocklist} def _update_models_from_file(self): """ 从 models.yaml 或覆盖配置文件中更新 MODEL_NAMES """ # 检查是否设置了覆盖配置文件的环境变量 override_config_path = os.getenv("OVERRIDE_DEFAULT_MODELS_CONFIG_WITH") if override_config_path and os.path.exists(override_config_path): config_file = Path(override_config_path) logger.info(f"Using override models config from: {override_config_path}") else: config_file = Path("src/config/static/models.yaml") logger.info("Using default models config") self._models_config_path = str(config_file) with open(self._models_config_path, encoding="utf-8") as f: _models = yaml.safe_load(f) self.model_names = _models["MODEL_NAMES"] self.embed_model_names = _models["EMBED_MODEL_INFO"] self.reranker_names = _models["RERANKER_LIST"] def _save_models_to_file(self): """ 将当前模型配置写回模型配置文件 """ if self._models_config_path is None: self._models_config_path = str(Path("src/config/static/models.yaml")) models_payload = { "MODEL_NAMES": self.model_names, "EMBED_MODEL_INFO": self.embed_model_names, "RERANKER_LIST": self.reranker_names, } with open(self._models_config_path, "w", encoding="utf-8") as f: yaml.safe_dump(models_payload, f, indent=2, allow_unicode=True, sort_keys=False) def handle_self(self): """ 处理配置 """ self.model_dir = os.environ.get("MODEL_DIR", "") if self.model_dir: if os.path.exists(self.model_dir): logger.debug( f"The model directory （{self.model_dir}） " f"contains the following folders: {os.listdir(self.model_dir)}" ) else: logger.warning( f"Warning: The model directory （{self.model_dir}） does not exist. " "If not configured, please ignore it. " "If configured, please check if the configuration is correct; " "For example, the mapping in the docker-compose file" ) # 检查模型提供商的环境变量 self.model_provider_status = {} for provider in self.model_names: env_var = self.model_names[provider]["env"] # 如果环境变量名为 NO_API_KEY，则认为总是可用 if env_var == "NO_API_KEY": self.model_provider_status[provider] = True else: self.model_provider_status[provider] = bool(os.getenv(env_var)) if os.getenv("TAVILY_API_KEY"): self.enable_web_search = True self.valuable_model_provider = [k for k, v in self.model_provider_status.items() if v] assert len(self.valuable_model_provider) > 0, "No model provider available, please check your `.env` file." def load(self): """根据传入的文件覆盖掉默认配置""" logger.info(f"Loading config from {self.filename}") if self.filename is not None and os.path.exists(self.filename): if self.filename.endswith(".json"): with open(self.filename) as f: content = f.read() if content: local_config = json.loads(content) self.update(local_config) else: print(f"{self.filename} is empty.") elif self.filename.endswith(".yaml"): with open(self.filename) as f: content = f.read() if content: local_config = yaml.safe_load(content) self.update(local_config) else: print(f"{self.filename} is empty.") else: logger.warning(f"Unknown config file type {self.filename}") def save(self): logger.info(f"Saving config to {self.filename}") if self.filename is None: logger.warning("Config file is not specified, save to default config/base.yaml") self.filename = os.path.join(self.save_dir, "config", "base.yaml") os.makedirs(os.path.dirname(self.filename), exist_ok=True) if self.filename.endswith(".json"): with open(self.filename, "w+") as f: json.dump(self.__dict__(), f, indent=4, ensure_ascii=False) elif self.filename.endswith(".yaml"): with open(self.filename, "w+") as f: yaml.dump(self.__dict__(), f, indent=2, allow_unicode=True) else: logger.warning(f"Unknown config file type {self.filename}, save as json") with open(self.filename, "w+") as f: json.dump(self, f, indent=4) logger.info(f"Config file {self.filename} saved") def dump_config(self): return json.loads(str(self)) config = Config()