scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 165

fix: address remaining greptile issues - Windows compat, YAML quoting, credential flow

Browse source 
- guard termios/tty imports with try/except ImportError for Windows compat
- quote master_key as YAML double-quoted scalar (same as env vars)
- remove unused port param from _build_config signature
- _validate_and_report now returns the final key so re-entered creds are stored
- add test for master_key YAML quoting
This commit is contained in:
Ishaan Jaffer 2026-03-14 15:33:46 -07:00
parent 6d95ddef4c
commit 1fbf098108
2 changed files with 53 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

47
litellm/setup_wizard.py
View file

@ -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:

35
tests/test_litellm/test_setup_wizard.py
View file

@ -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