"""Shared configuration helpers for CLI and UI.""" from __future__ import annotations import re from typing import Any, Dict, Optional from chemgraph.models.supported_models import ( ALCF_DEFAULT_BASE_URL, ARGO_DEFAULT_BASE_URL, all_supported_models, supported_alcf_models, supported_anthropic_models, supported_argo_models, supported_gemini_models, supported_ollama_models, supported_openai_models, ) def flatten_config(config: Dict[str, Any]) -> Dict[str, Any]: """Flatten nested TOML-like config into top-level keys used by the CLI. Parameters ---------- config : dict[str, Any] Nested configuration dictionary. Returns ------- dict[str, Any] Flattened configuration with section names included in keys. """ flattened: Dict[str, Any] = {} if "general" in config: flattened.update(config["general"]) for section in ["api", "chemistry", "output"]: if section in config: for key, value in config[section].items(): if isinstance(value, dict): for subkey, subvalue in value.items(): flattened[f"{section}_{key}_{subkey}"] = subvalue else: flattened[f"{section}_{key}"] = value for section in ["logging", "features", "security", "advanced"]: if section in config: if isinstance(config[section], dict): for key, value in config[section].items(): if isinstance(value, dict): for subkey, subvalue in value.items(): flattened[f"{section}_{key}_{subkey}"] = subvalue else: flattened[f"{section}_{key}"] = value else: flattened[section] = config[section] return flattened def normalize_openai_base_url(base_url: Optional[str]) -> Optional[str]: """Normalize Argo-style URLs to OpenAI-compatible /v1 URLs. Parameters ---------- base_url : str, optional Provider base URL. Returns ------- str or None Normalized URL, or ``None`` when no URL was provided. """ if not base_url: return base_url if ( "apps-dev.inside.anl.gov/argoapi" in base_url or "apps.inside.anl.gov/argoapi" in base_url ): base_url = re.sub(r"/api/v1/resource/(chat|embed)/?$", "/v1", base_url) base_url = re.sub(r"/docs/?$", "", base_url) base_url = re.sub(r"/api/v1/?$", "/v1", base_url) base_url = base_url.rstrip("/") return base_url def get_base_url_for_model_from_nested_config( model_name: str, config: Dict[str, Any] ) -> Optional[str]: """Resolve provider base URL using nested config structure. Parameters ---------- model_name : str Model identifier. config : dict[str, Any] Nested configuration dictionary. Returns ------- str or None Matching provider base URL, or ``None`` when not configured. """ api = config.get("api", {}) if model_name in supported_argo_models: return normalize_openai_base_url( api.get("openai", {}).get("base_url") or ARGO_DEFAULT_BASE_URL ) if model_name in supported_openai_models: return normalize_openai_base_url(api.get("openai", {}).get("base_url")) if model_name in supported_alcf_models: return api.get("alcf", {}).get("base_url") or ALCF_DEFAULT_BASE_URL if model_name in supported_anthropic_models: return api.get("anthropic", {}).get("base_url") if model_name in supported_gemini_models: return api.get("google", {}).get("base_url") if model_name in supported_ollama_models: return api.get("local", {}).get("base_url") return normalize_openai_base_url(api.get("openai", {}).get("base_url")) def get_base_url_for_model_from_flat_config( model_name: str, config: Dict[str, Any] ) -> Optional[str]: """Resolve provider base URL using flattened config keys. Parameters ---------- model_name : str Model identifier. config : dict[str, Any] Flattened configuration dictionary. Returns ------- str or None Matching provider base URL, or ``None`` when not configured. """ if model_name in supported_argo_models: return normalize_openai_base_url( config.get("api_openai_base_url") or ARGO_DEFAULT_BASE_URL ) if model_name in supported_openai_models: return normalize_openai_base_url(config.get("api_openai_base_url")) if model_name in supported_alcf_models: return config.get("api_alcf_base_url") or ALCF_DEFAULT_BASE_URL if model_name in supported_anthropic_models: return config.get("api_anthropic_base_url") if model_name in supported_gemini_models: return config.get("api_google_base_url") if model_name in supported_ollama_models: return config.get("api_local_base_url") return normalize_openai_base_url(config.get("api_openai_base_url")) def get_model_options_for_nested_config(config: Dict[str, Any]) -> list[str]: """Return model options for UI selection. Always show all curated models so users can switch providers from the UI. If Argo endpoint is configured, prioritize Argo model IDs at the top. Parameters ---------- config : dict[str, Any] Nested configuration dictionary. Returns ------- list[str] Model identifiers for UI selection. """ base_url = config.get("api", {}).get("openai", {}).get("base_url") if base_url and "argoapi" in base_url: remaining = [m for m in all_supported_models if m not in supported_argo_models] return supported_argo_models + remaining return all_supported_models def get_argo_user_from_nested_config(config: Dict[str, Any]) -> Optional[str]: """Resolve Argo user from nested config. Parameters ---------- config : dict[str, Any] Nested configuration dictionary. Returns ------- str or None Configured Argo username, or ``None``. """ value = config.get("api", {}).get("openai", {}).get("argo_user") if isinstance(value, str): value = value.strip() return value or None def get_argo_user_from_flat_config(config: Dict[str, Any]) -> Optional[str]: """Resolve Argo user from flattened config. Parameters ---------- config : dict[str, Any] Flattened configuration dictionary. Returns ------- str or None Configured Argo username, or ``None``. """ value = config.get("api_openai_argo_user") if isinstance(value, str): value = value.strip() return value or None