import ast import os import sys from typing import Any, Literal, Optional import orjson from loguru import logger from pydantic import BaseModel, Field, ValidationError, field_validator from pydantic_settings import ( BaseSettings, SettingsConfigDict, YamlConfigSettingsSource, ) CONFIG_PATH = "config/config.yaml" class HTTPSConfig(BaseModel): """HTTPS configuration""" enabled: bool = Field(default=False, description="Enable HTTPS") key_file: str = Field(default="certs/privkey.pem", description="SSL private key file path") cert_file: str = Field(default="certs/fullchain.pem", description="SSL certificate file path") class ServerConfig(BaseModel): """Server configuration""" host: str = Field(default="0.0.0.0", description="Server host address") port: int = Field(default=8000, ge=1, le=65535, description="Server port number") api_key: Optional[str] = Field( default=None, description="API key for authentication, if set, will enable API key validation", ) https: HTTPSConfig = Field(default=HTTPSConfig(), description="HTTPS configuration") class GeminiClientSettings(BaseModel): """Credential set for one Gemini client.""" id: str = Field(..., description="Unique identifier for the client") secure_1psid: str = Field(..., description="Gemini Secure 1PSID") secure_1psidts: str = Field(..., description="Gemini Secure 1PSIDTS") proxy: Optional[str] = Field(default=None, description="Proxy URL for this Gemini client") @field_validator("proxy", mode="before") @classmethod def _blank_proxy_to_none(cls, value: Optional[str]) -> Optional[str]: if value is None: return None stripped = value.strip() return stripped or None class GeminiModelConfig(BaseModel): """Configuration for a custom Gemini model.""" model_name: Optional[str] = Field(default=None, description="Name of the model") model_header: Optional[dict[str, Optional[str]]] = Field( default=None, description="Header for the model" ) @field_validator("model_header", mode="before") @classmethod def _parse_json_string(cls, v: Any) -> Any: if isinstance(v, str) and v.strip().startswith("{"): try: return orjson.loads(v) except orjson.JSONDecodeError: # Return the original value to let Pydantic handle the error or type mismatch return v return v class GeminiConfig(BaseModel): """Gemini API configuration""" clients: list[GeminiClientSettings] = Field( ..., description="List of Gemini client credential pairs" ) models: list[GeminiModelConfig] = Field(default=[], description="List of custom Gemini models") model_strategy: Literal["append", "overwrite"] = Field( default="append", description="Strategy for loading models: 'append' merges custom with default, 'overwrite' uses only custom", ) timeout: int = Field(default=120, ge=1, description="Init timeout") auto_refresh: bool = Field(True, description="Enable auto-refresh for Gemini cookies") refresh_interval: int = Field( default=540, ge=1, description="Interval in seconds to refresh Gemini cookies" ) verbose: bool = Field(False, description="Enable verbose logging for Gemini API requests") max_chars_per_request: int = Field( default=1_000_000, ge=1, description="Maximum characters Gemini Web can accept per request", ) @field_validator("models", mode="before") @classmethod def _parse_models_json(cls, v: Any) -> Any: if isinstance(v, str) and v.strip().startswith("["): try: return orjson.loads(v) except orjson.JSONDecodeError as e: logger.warning(f"Failed to parse models JSON string: {e}") return v return v @field_validator("models") @classmethod def _filter_valid_models(cls, v: list[GeminiModelConfig]) -> list[GeminiModelConfig]: """Filter out models that don't have all required fields set.""" valid_models = [] for model in v: if model.model_name and model.model_header: valid_models.append(model) else: missing = [] if not model.model_name: missing.append("model_name") if not model.model_header: missing.append("model_header") logger.warning( f"Discarding custom model due to missing {', '.join(missing)}: {model}" ) return valid_models class CORSConfig(BaseModel): """CORS configuration""" enabled: bool = Field(default=True, description="Enable CORS support") allow_origins: list[str] = Field( default=["*"], description="List of allowed origins for CORS requests" ) allow_credentials: bool = Field(default=True, description="Allow credentials in CORS requests") allow_methods: list[str] = Field( default=["*"], description="List of allowed HTTP methods for CORS requests" ) allow_headers: list[str] = Field( default=["*"], description="List of allowed headers for CORS requests" ) class StorageConfig(BaseModel): """LMDB Storage configuration""" path: str = Field( default="data/lmdb", description="Path to the storage directory where data will be saved", ) max_size: int = Field( default=1024**2 * 256, # 256 MB ge=1, description="Maximum size of the storage in bytes", ) retention_days: int = Field( default=14, ge=0, description="Number of days to retain conversations before automatic cleanup (0 disables cleanup)", ) class LoggingConfig(BaseModel): """Logging configuration""" level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = Field( default="DEBUG", description="Logging level", ) class Config(BaseSettings): """Application configuration""" # Server configuration server: ServerConfig = Field( default=ServerConfig(), description="Server configuration, including host, port, and API key", ) # CORS configuration cors: CORSConfig = Field( default=CORSConfig(), description="CORS configuration, allows cross-origin requests", ) # Gemini API configuration gemini: GeminiConfig = Field(..., description="Gemini API configuration, must be set") storage: StorageConfig = Field( default=StorageConfig(), description="Storage configuration, defines where and how data will be stored", ) # Logging configuration logging: LoggingConfig = Field( default=LoggingConfig(), description="Logging configuration", ) model_config = SettingsConfigDict( env_prefix="CONFIG_", env_nested_delimiter="__", nested_model_default_partial_update=True, yaml_file=os.getenv("CONFIG_PATH", CONFIG_PATH), ) @classmethod def settings_customise_sources( cls, settings_cls, init_settings, env_settings, dotenv_settings, file_secret_settings, ): """Read settings: env -> yaml -> default""" return ( env_settings, YamlConfigSettingsSource(settings_cls), ) def extract_gemini_clients_env() -> dict[int, dict[str, str]]: """Extract and remove all Gemini clients related environment variables, return a mapping from index to field dict.""" prefix = "CONFIG_GEMINI__CLIENTS__" env_overrides: dict[int, dict[str, str]] = {} to_delete = [] for k, v in os.environ.items(): if k.startswith(prefix): parts = k.split("__") if len(parts) < 4: continue index_str, field = parts[2], parts[3].lower() if not index_str.isdigit(): continue idx = int(index_str) env_overrides.setdefault(idx, {})[field] = v to_delete.append(k) # Remove these environment variables to avoid Pydantic parsing errors for k in to_delete: del os.environ[k] return env_overrides def _merge_clients_with_env( base_clients: list[GeminiClientSettings] | None, env_overrides: dict[int, dict[str, str]], ): """Override base_clients with env_overrides, return the new clients list.""" if not env_overrides: return base_clients result_clients: list[GeminiClientSettings] = [] if base_clients: result_clients = [client.model_copy() for client in base_clients] for idx in sorted(env_overrides): overrides = env_overrides[idx] if idx < len(result_clients): client_dict = result_clients[idx].model_dump() client_dict.update(overrides) result_clients[idx] = GeminiClientSettings(**client_dict) elif idx == len(result_clients): new_client = GeminiClientSettings(**overrides) result_clients.append(new_client) else: raise IndexError( f"Client index {idx} in env is out of range (current count: {len(result_clients)}). " "Client indices must be contiguous starting from 0." ) return result_clients if result_clients else base_clients def extract_gemini_models_env() -> dict[int, dict[str, Any]]: """Extract and remove all Gemini models related environment variables, supporting nested fields.""" root_key = "CONFIG_GEMINI__MODELS" env_overrides: dict[int, dict[str, Any]] = {} if root_key in os.environ: val = os.environ[root_key] models_list = None parsed_successfully = False try: models_list = orjson.loads(val) parsed_successfully = True except orjson.JSONDecodeError: try: models_list = ast.literal_eval(val) parsed_successfully = True except (ValueError, SyntaxError) as e: logger.warning(f"Failed to parse {root_key} as JSON or Python literal: {e}") if parsed_successfully and isinstance(models_list, list): for idx, model_data in enumerate(models_list): if isinstance(model_data, dict): env_overrides[idx] = model_data # Remove the environment variable to avoid Pydantic parsing errors del os.environ[root_key] return env_overrides def _merge_models_with_env( base_models: list[GeminiModelConfig] | None, env_overrides: dict[int, dict[str, Any]], ): """Override base_models with env_overrides using standard update (replace whole fields).""" if not env_overrides: return base_models or [] result_models: list[GeminiModelConfig] = [] if base_models: result_models = [model.model_copy() for model in base_models] for idx in sorted(env_overrides): overrides = env_overrides[idx] if idx < len(result_models): # Update existing model: overwrite fields found in env model_dict = result_models[idx].model_dump() model_dict.update(overrides) result_models[idx] = GeminiModelConfig(**model_dict) elif idx == len(result_models): # Append new models new_model = GeminiModelConfig(**overrides) result_models.append(new_model) else: raise IndexError( f"Model index {idx} in env is out of range (current count: {len(result_models)}). " "Model indices must be contiguous starting from 0." ) return result_models def initialize_config() -> Config: """ Initialize the configuration. Returns: Config: Configuration object """ try: # First, extract and remove Gemini clients related environment variables env_clients_overrides = extract_gemini_clients_env() # Extract and remove Gemini models related environment variables env_models_overrides = extract_gemini_models_env() # Then, initialize Config with pydantic_settings config = Config() # type: ignore # Synthesize clients config.gemini.clients = _merge_clients_with_env( config.gemini.clients, env_clients_overrides ) # Synthesize models config.gemini.models = _merge_models_with_env(config.gemini.models, env_models_overrides) return config except ValidationError as e: logger.error(f"Configuration validation failed: {e!s}") sys.exit(1)