diff --git a/litellm/setup_wizard.py b/litellm/setup_wizard.py index 822dcb31cdc..3a579d727c3 100644 --- a/litellm/setup_wizard.py +++ b/litellm/setup_wizard.py @@ -13,11 +13,20 @@ import re import secrets import sys import sysconfig -import termios -import tty from pathlib import Path from typing import Dict, List, Optional, Set +# termios / tty are Unix-only; fall back gracefully on Windows +try: + import termios + import tty + + _HAS_RAW_TERMINAL: bool = True +except ImportError: + termios = None # type: ignore[assignment] + tty = None # type: ignore[assignment] + _HAS_RAW_TERMINAL = False + from litellm.utils import check_valid_key # --------------------------------------------------------------------------- @@ -221,7 +230,7 @@ class SetupWizard: config_path = Path(os.getcwd()) / "litellm_config.yaml" try: config_path.write_text( - SetupWizard._build_config(providers, env_vars, port, master_key) + SetupWizard._build_config(providers, env_vars, master_key) ) except OSError as exc: print(f"\n {bold(_CROSS + ' Could not write config:')} {exc}") @@ -251,14 +260,19 @@ class SetupWizard: @staticmethod def _select_providers() -> List[Dict]: """Arrow-key multi-select. Falls back to number input if /dev/tty unavailable.""" + if not _HAS_RAW_TERMINAL: + return SetupWizard._select_fallback() try: return SetupWizard._select_interactive() - except (OSError, termios.error): + except OSError: return SetupWizard._select_fallback() @staticmethod def _read_key() -> str: """Read one keypress from /dev/tty in raw mode.""" + assert ( + termios is not None and tty is not None + ) # only called when _HAS_RAW_TERMINAL with open("/dev/tty", "rb") as tty_fh: fd = tty_fh.fileno() old = termios.tcgetattr(fd) @@ -390,8 +404,6 @@ class SetupWizard: if not key: continue - env_vars[p["env_key"]] = key - for extra_key, extra_hint in zip( p.get("extra_keys", []), p.get("extra_hints", []) ): @@ -413,7 +425,8 @@ class SetupWizard: deployment ) - SetupWizard._validate_and_report(p, key) + # Store the key returned by validation — may be a re-entered replacement + env_vars[p["env_key"]] = SetupWizard._validate_and_report(p, key) return env_vars @@ -432,14 +445,14 @@ class SetupWizard: return "" @staticmethod - def _validate_and_report(provider: Dict, api_key: str) -> None: + def _validate_and_report(provider: Dict, api_key: str) -> str: """ Validate credentials using litellm.utils.check_valid_key and print result. - Offers a re-entry loop on failure. + Offers a re-entry loop on failure. Returns the final (possibly re-entered) key. """ test_model: Optional[str] = provider.get("test_model") if not test_model: - return # Azure / Bedrock / Ollama — skip + return api_key # Azure / Bedrock / Ollama — skip validation while True: print( @@ -451,21 +464,21 @@ class SetupWizard: print( f" {green(_CHECK)} {bold(provider['name'])} connected successfully" ) - return + return api_key print(f" {_CROSS} {bold(provider['name'])} {grey('— invalid API key')}") if ( _styled_input(f" {blue('❯')} Re-enter key? {grey('(y/N)')}: ").lower() != "y" ): - return + return api_key hint = grey(provider.get("key_hint", "")) new_key = _styled_input( f" {blue('❯')} {bold(provider['name'])} API key {hint}: " ) if not new_key: - return + return api_key api_key = new_key # ── proxy settings ─────────────────────────────────────────────────────── @@ -496,7 +509,6 @@ class SetupWizard: def _build_config( providers: List[Dict], env_vars: Dict[str, str], - port: int, master_key: str, ) -> str: env_copy = dict(env_vars) # work on a copy — do not mutate caller's dict @@ -535,7 +547,12 @@ class SetupWizard: if p.get("api_version"): lines.append(f" api_version: {p['api_version']}") - lines += ["", "general_settings:", f" master_key: {master_key}", ""] + lines += [ + "", + "general_settings:", + f' master_key: "{_yaml_escape(master_key)}"', + "", + ] real_vars = {k: v for k, v in env_copy.items() if not k.startswith("_LITELLM_")} if real_vars: diff --git a/tests/test_litellm/test_setup_wizard.py b/tests/test_litellm/test_setup_wizard.py index 545fe40e70c..49a1711b0dc 100644 --- a/tests/test_litellm/test_setup_wizard.py +++ b/tests/test_litellm/test_setup_wizard.py @@ -68,14 +68,13 @@ def test_build_config_basic_openai(): config = SetupWizard._build_config( [_OPENAI], {"OPENAI_API_KEY": "sk-test"}, - 4000, "sk-master", ) assert "model_list:" in config assert "model_name: gpt-4o" in config assert "model: gpt-4o" in config assert "api_key: os.environ/OPENAI_API_KEY" in config - assert "master_key: sk-master" in config + assert 'master_key: "sk-master"' in config def test_build_config_skipped_provider_omitted(): @@ -83,7 +82,6 @@ def test_build_config_skipped_provider_omitted(): config = SetupWizard._build_config( [_OPENAI, _ANTHROPIC], {"ANTHROPIC_API_KEY": "sk-ant-test"}, # OpenAI key missing - 4000, "sk-master", ) assert "gpt-4o" not in config @@ -95,12 +93,21 @@ def test_build_config_env_vars_written_escaped(): config = SetupWizard._build_config( [_OPENAI], {"OPENAI_API_KEY": 'sk-ab"cd'}, - 4000, "sk-master", ) assert 'OPENAI_API_KEY: "sk-ab\\"cd"' in config +def test_build_config_master_key_quoted(): + """master_key must be quoted in YAML to handle special characters.""" + config = SetupWizard._build_config( + [_OPENAI], + {"OPENAI_API_KEY": "sk-test"}, + 'sk-master"special', + ) + assert 'master_key: "sk-master\\"special"' in config + + def test_build_config_does_not_mutate_env_vars(): """_build_config must not modify the caller's env_vars dict.""" env_vars = { @@ -109,7 +116,7 @@ def test_build_config_does_not_mutate_env_vars(): "_LITELLM_AZURE_DEPLOYMENT_AZURE": "my-deployment", } original_keys = set(env_vars.keys()) - SetupWizard._build_config([_AZURE], env_vars, 4000, "sk-master") + SetupWizard._build_config([_AZURE], env_vars, "sk-master") assert set(env_vars.keys()) == original_keys @@ -119,7 +126,7 @@ def test_build_config_azure_uses_deployment_name(): "_LITELLM_AZURE_API_BASE_AZURE": "https://my.azure.com", "_LITELLM_AZURE_DEPLOYMENT_AZURE": "my-gpt4o", } - config = SetupWizard._build_config([_AZURE], env_vars, 4000, "sk-master") + config = SetupWizard._build_config([_AZURE], env_vars, "sk-master") assert "model: azure/my-gpt4o" in config assert "model_name: azure-my-gpt4o" in config @@ -131,22 +138,22 @@ def test_build_config_no_display_name_collision_openai_and_azure(): "AZURE_API_KEY": "az-key", "_LITELLM_AZURE_DEPLOYMENT_AZURE": "gpt-4o", } - config = SetupWizard._build_config([_OPENAI, _AZURE], env_vars, 4000, "sk-master") - assert "model_name: gpt-4o" in config # OpenAI + config = SetupWizard._build_config([_OPENAI, _AZURE], env_vars, "sk-master") + assert "model_name: gpt-4o" in config # OpenAI assert "model_name: azure-gpt-4o" in config # Azure — qualified def test_build_config_ollama_no_api_key_line(): """Ollama has no env_key — config should not contain an api_key line for it.""" - config = SetupWizard._build_config([_OLLAMA], {}, 4000, "sk-master") + config = SetupWizard._build_config([_OLLAMA], {}, "sk-master") assert "ollama/llama3.2" in config assert "api_key:" not in config -def test_build_config_port_in_general_settings(): - """Port is not currently written to general_settings, but master_key is.""" - config = SetupWizard._build_config([_OPENAI], {"OPENAI_API_KEY": "k"}, 8080, "sk-m") - assert "master_key: sk-m" in config +def test_build_config_master_key_in_general_settings(): + """master_key is written to general_settings.""" + config = SetupWizard._build_config([_OPENAI], {"OPENAI_API_KEY": "k"}, "sk-m") + assert 'master_key: "sk-m"' in config def test_build_config_internal_sentinel_keys_excluded(): @@ -155,5 +162,5 @@ def test_build_config_internal_sentinel_keys_excluded(): "OPENAI_API_KEY": "sk-real", "_LITELLM_AZURE_API_BASE_AZURE": "https://x.azure.com", } - config = SetupWizard._build_config([_OPENAI], env_vars, 4000, "sk-master") + config = SetupWizard._build_config([_OPENAI], env_vars, "sk-master") assert "_LITELLM_" not in config