"""Application configuration.""" from __future__ import annotations import os from pathlib import Path from typing import Any import tomli import tomli_w from pydantic import BaseModel, Field, PrivateAttr from src.config.static.models import ( DEFAULT_CHAT_MODEL_PROVIDERS, DEFAULT_EMBED_MODELS, DEFAULT_RERANKERS, ChatModelProvider, EmbedModelInfo, RerankerInfo, ) from src.utils.logging_config import logger class Config(BaseModel): save_dir: str = Field(default="saves", description="Storage root directory") model_dir: str = Field(default="", description="Local model directory") enable_reranker: bool = Field(default=False) enable_content_guard: bool = Field(default=False) enable_content_guard_llm: bool = Field(default=False) enable_web_search: bool = Field(default=False) default_model: str = Field(default="siliconflow/deepseek-ai/DeepSeek-V3.2") fast_model: str = Field(default="siliconflow/THUDM/GLM-4-9B-0414") embed_model: str = Field(default="siliconflow/BAAI/bge-m3") reranker: str = Field(default="siliconflow/BAAI/bge-reranker-v2-m3") content_guard_llm_model: str = Field(default="siliconflow/Qwen/Qwen3-235B-A22B-Instruct-2507") default_agent_id: str = Field(default="") sandbox_provider: str = Field(default="provisioner") sandbox_provisioner_url: str = Field(default="http://sandbox-provisioner:8002") sandbox_virtual_path_prefix: str = Field(default="/mnt/user-data") sandbox_exec_timeout_seconds: int = Field(default=180) sandbox_max_output_bytes: int = Field(default=262144) model_names: dict[str, ChatModelProvider] = Field( default_factory=lambda: DEFAULT_CHAT_MODEL_PROVIDERS.copy(), exclude=True, ) embed_model_names: dict[str, EmbedModelInfo] = Field( default_factory=lambda: DEFAULT_EMBED_MODELS.copy(), exclude=True, ) reranker_names: dict[str, RerankerInfo] = Field( default_factory=lambda: DEFAULT_RERANKERS.copy(), exclude=True, ) model_provider_status: dict[str, bool] = Field(default_factory=dict, exclude=True) valuable_model_provider: list[str] = Field(default_factory=list, exclude=True) _config_file: Path | None = PrivateAttr(default=None) _user_modified_fields: set[str] = PrivateAttr(default_factory=set) _modified_providers: set[str] = PrivateAttr(default_factory=set) model_config = {"arbitrary_types_allowed": True, "extra": "allow"} def __init__(self, **data): super().__init__(**data) self._setup_paths() self._load_user_config() self._load_custom_providers() self._handle_environment() def _setup_paths(self) -> None: self.save_dir = os.getenv("SAVE_DIR") or self.save_dir self._config_file = Path(self.save_dir) / "config" / "base.toml" self._config_file.parent.mkdir(parents=True, exist_ok=True) def _load_user_config(self) -> None: if not self._config_file or not self._config_file.exists(): logger.info(f"Config file not found, using defaults: {self._config_file}") return try: with self._config_file.open("rb") as file: user_config = tomli.load(file) except Exception as exc: # noqa: BLE001 logger.error(f"Failed to load config from {self._config_file}: {exc}") return self._user_modified_fields = set(user_config.keys()) for key, value in user_config.items(): if key == "model_names": self._load_model_names(value) elif key in self.model_fields: setattr(self, key, value) else: logger.warning(f"Unknown config key: {key}") def _load_model_names(self, model_names_data: dict[str, Any]) -> None: for provider, provider_data in (model_names_data or {}).items(): try: if provider in self.model_names: merged = self.model_names[provider].model_dump() | dict(provider_data or {}) self.model_names[provider] = ChatModelProvider(**merged) else: self.model_names[provider] = ChatModelProvider(**provider_data) except Exception as exc: # noqa: BLE001 logger.warning(f"Skip invalid model provider config {provider}: {exc}") def _load_custom_providers(self) -> None: if not self._config_file: return custom_config_file = self._config_file.parent / "custom_providers.toml" if not custom_config_file.exists(): return try: with custom_config_file.open("rb") as file: custom_config = tomli.load(file) except Exception as exc: # noqa: BLE001 logger.error(f"Failed to load custom providers from {custom_config_file}: {exc}") return model_names = custom_config.get("model_names") or {} self._load_custom_model_providers(model_names) def _load_custom_model_providers(self, providers_data: dict[str, Any]) -> None: for provider, provider_data in (providers_data or {}).items(): try: payload = dict(provider_data or {}) payload["custom"] = True self.model_names[provider] = ChatModelProvider(**payload) except Exception as exc: # noqa: BLE001 logger.warning(f"Skip invalid custom provider {provider}: {exc}") def _handle_environment(self) -> None: self.model_dir = os.environ.get("MODEL_DIR") or self.model_dir self.model_provider_status = {} for provider, info in self.model_names.items(): env_var = (info.env or "").strip() if env_var.upper() == "NO_API_KEY": self.model_provider_status[provider] = True continue api_key = os.environ.get(env_var) self.model_provider_status[provider] = bool(api_key or info.custom) if os.getenv("TAVILY_API_KEY"): self.enable_web_search = True self.valuable_model_provider = [key for key, ok in self.model_provider_status.items() if ok] self.sandbox_provider = (os.getenv("SANDBOX_PROVIDER") or self.sandbox_provider or "provisioner").strip() self.sandbox_provisioner_url = ( os.getenv("SANDBOX_PROVISIONER_URL") or self.sandbox_provisioner_url or "http://sandbox-provisioner:8002" ).strip() self.sandbox_virtual_path_prefix = ( os.getenv("SANDBOX_VIRTUAL_PATH_PREFIX") or self.sandbox_virtual_path_prefix or "/mnt/user-data" ).strip() self.sandbox_exec_timeout_seconds = int( os.getenv("SANDBOX_EXEC_TIMEOUT_SECONDS") or self.sandbox_exec_timeout_seconds or 180 ) self.sandbox_max_output_bytes = int( os.getenv("SANDBOX_MAX_OUTPUT_BYTES") or self.sandbox_max_output_bytes or 262144 ) if self.sandbox_provider.lower() != "provisioner": raise ValueError("Only sandbox_provider=provisioner is supported.") if not self.sandbox_provisioner_url: raise ValueError("SANDBOX_PROVISIONER_URL is required when sandbox provider is provisioner.") if not self.sandbox_virtual_path_prefix.startswith("/"): self.sandbox_virtual_path_prefix = f"/{self.sandbox_virtual_path_prefix}" if not self.valuable_model_provider: raise ValueError("No model provider available, please check your `.env` file.") def save(self) -> None: if not self._config_file: logger.warning("Config file path not set") return default_config = Config.model_construct() user_modified: dict[str, Any] = {} for field_name, field_info in self.model_fields.items(): if field_info.exclude: continue current_value = getattr(self, field_name) default_value = getattr(default_config, field_name) if current_value != default_value: user_modified[field_name] = current_value try: with self._config_file.open("wb") as file: tomli_w.dump(user_modified, file) except Exception as exc: # noqa: BLE001 logger.error(f"Failed to save config to {self._config_file}: {exc}") def dump_config(self) -> dict[str, Any]: config_dict = self.model_dump( exclude={ "model_names", "embed_model_names", "reranker_names", "model_provider_status", "valuable_model_provider", } ) config_dict["model_names"] = {provider: info.model_dump() for provider, info in self.model_names.items()} config_dict["embed_model_names"] = { model_id: info.model_dump() for model_id, info in self.embed_model_names.items() } config_dict["reranker_names"] = {model_id: info.model_dump() for model_id, info in self.reranker_names.items()} config_dict["model_provider_status"] = self.model_provider_status config_dict["valuable_model_provider"] = self.valuable_model_provider fields_info: dict[str, Any] = {} for field_name, field_info in Config.model_fields.items(): if field_info.exclude: continue annotation = field_info.annotation fields_info[field_name] = { "des": field_info.description, "default": field_info.default, "type": annotation.__name__ if hasattr(annotation, "__name__") else str(annotation), "exclude": bool(field_info.exclude), } config_dict["_config_items"] = fields_info return config_dict def get_model_choices(self) -> list[str]: choices: list[str] = [] for provider, info in self.model_names.items(): if not self.model_provider_status.get(provider, False): continue choices.extend([f"{provider}/{model}" for model in info.models]) return choices def get_embed_model_choices(self) -> list[str]: return list(self.embed_model_names.keys()) def get_reranker_choices(self) -> list[str]: return list(self.reranker_names.keys()) def __getitem__(self, key: str) -> Any: logger.warning("Using deprecated dict-style access for Config. Please use attribute access instead.") return getattr(self, key, None) def __setitem__(self, key: str, value: Any) -> None: logger.warning("Using deprecated dict-style assignment for Config. Please use attribute access instead.") setattr(self, key, value) def update(self, other: dict[str, Any]) -> None: for key, value in (other or {}).items(): if key in self.model_fields: setattr(self, key, value) else: logger.warning(f"Unknown config key: {key}") def _save_models_to_file(self, provider_name: str | None = None) -> None: if not self._config_file: logger.warning("Config file path not set") return user_config: dict[str, Any] = {} if self._config_file.exists(): with self._config_file.open("rb") as file: user_config = tomli.load(file) user_config.setdefault("model_names", {}) if provider_name: if provider_name in self.model_names: user_config["model_names"][provider_name] = self.model_names[provider_name].model_dump() self._modified_providers.add(provider_name) else: user_config["model_names"] = {provider: info.model_dump() for provider, info in self.model_names.items()} self._user_modified_fields.add("model_names") with self._config_file.open("wb") as file: tomli_w.dump(user_config, file) def get_custom_providers(self) -> dict[str, ChatModelProvider]: return {provider: info for provider, info in self.model_names.items() if info.custom} def _save_custom_providers(self) -> None: if not self._config_file: logger.warning("Config file path not set") return custom_config_file = self._config_file.parent / "custom_providers.toml" custom_providers = self.get_custom_providers() custom_config: dict[str, Any] = {} if custom_providers: custom_config["model_names"] = {provider: info.model_dump() for provider, info in custom_providers.items()} custom_config_file.parent.mkdir(parents=True, exist_ok=True) with custom_config_file.open("wb") as file: tomli_w.dump(custom_config, file) def add_custom_provider(self, provider_id: str, provider_data: dict[str, Any]) -> bool: if provider_id in self.model_names: logger.error(f"Provider ID already exists: {provider_id}") return False payload = dict(provider_data or {}) env_value = payload.get("env") if isinstance(env_value, str) and env_value.startswith("${") and env_value.endswith("}"): payload["env"] = env_value[2:-1] payload["custom"] = True self.model_names[provider_id] = ChatModelProvider(**payload) self._save_custom_providers() self._handle_environment() return True def update_custom_provider(self, provider_id: str, provider_data: dict[str, Any]) -> bool: if provider_id not in self.model_names: logger.error(f"Provider not found: {provider_id}") return False if not self.model_names[provider_id].custom: logger.error(f"Cannot update non-custom provider: {provider_id}") return False payload = dict(provider_data or {}) env_value = payload.get("env") if isinstance(env_value, str) and env_value.startswith("${") and env_value.endswith("}"): payload["env"] = env_value[2:-1] payload["custom"] = True self.model_names[provider_id] = ChatModelProvider(**payload) self._save_custom_providers() self._handle_environment() return True def delete_custom_provider(self, provider_id: str) -> bool: if provider_id not in self.model_names: logger.error(f"Provider not found: {provider_id}") return False if not self.model_names[provider_id].custom: logger.error(f"Cannot delete non-custom provider: {provider_id}") return False del self.model_names[provider_id] self._save_custom_providers() self._handle_environment() return True config = Config()