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

address greptile review: fix yaml escaping, port validation, display name collisions, tests

Browse source 
- setup_wizard.py: add _yaml_escape() for safe YAML embedding of API keys
- setup_wizard.py: add _styled_input() with readline ANSI ignore markers
- setup_wizard.py: change DIVIDER to _divider() fn to avoid import-time color capture
- setup_wizard.py: validate port range 1-65535, initialize before loop
- setup_wizard.py: qualify azure display names (azure-gpt-4o) to avoid collision with openai
- setup_wizard.py: work on env_copy in _build_config to avoid mutating caller's dict
- setup_wizard.py: skip model_list entries for providers with no credentials
- setup_wizard.py: prompt for azure deployment name
- setup_wizard.py: wrap os.execlp in try/except with friendly fallback
- setup_wizard.py: wrap config write in try/except OSError
- setup_wizard.py: fix _validate_and_report to use two print lines (no \r overwrite)
- setup_wizard.py: add .gitignore tip next to key storage notice
- setup_wizard.py: fix run_setup_wizard() return type annotation to None
- scripts/install.sh: drop pipefail (not supported by dash on Ubuntu when invoked as sh)
- scripts/install.sh: use litellm[proxy] from PyPI (not hardcoded dev branch)
- scripts/install.sh: guard /dev/tty read with -r check for Docker/CI compat
- scripts/install.sh: remove --force-reinstall to avoid downgrading dependencies
- tests/test_litellm/test_setup_wizard.py: 13 unit tests for _build_config and _yaml_escape
This commit is contained in:
Ishaan Jaffer 2026-03-14 15:16:11 -07:00
parent cd85a3b2c7
commit 7ab9d06c3e
3 changed files with 277 additions and 49 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

148
litellm/setup_wizard.py
View file

@ -9,6 +9,7 @@ and generating a proxy config file — mirroring the Claude Code onboarding UX.
import importlib.metadata
import os
import re
import secrets
import sys
import sysconfig
@ -97,6 +98,8 @@ PROVIDERS: List[Dict] = [
# ANSI colour helpers
# ---------------------------------------------------------------------------
_ANSI_RE = re.compile(r"\033\[[^m]*m")
_ORANGE = "\033[38;2;215;119;87m"
_DIM = "\033[2m"
_BOLD = "\033[1m"
@ -144,6 +147,29 @@ def dim(t: str) -> str:
return _c(_DIM, t)
def _divider() -> str:
"""Return a styled divider line (evaluated at call-time, not import-time)."""
return dim(" " + "╌" * 74)
def _styled_input(prompt: str) -> str:
"""
Like input() but wraps ANSI sequences in readline ignore markers
(\\001...\\002) so readline correctly tracks the cursor column.
In non-TTY contexts, strips ANSI entirely so no escape codes appear.
"""
if sys.stdout.isatty():
rl_prompt = _ANSI_RE.sub(lambda m: f"\001{m.group()}\002", prompt)
else:
rl_prompt = _ANSI_RE.sub("", prompt)
return input(rl_prompt).strip()
def _yaml_escape(value: str) -> str:
"""Escape a string for safe embedding in a double-quoted YAML scalar."""
return value.replace("\\", "\\\\").replace('"', '\\"')
# ---------------------------------------------------------------------------
# Layout constants
# ---------------------------------------------------------------------------
@ -157,13 +183,12 @@ LITELLM_ASCII = r"""
╚══════╝╚═╝ ╚═╝ ╚══════╝╚══════╝╚══════╝╚═╝ ╚═╝
"""
DIVIDER = dim(" " + "╌" * 74)
# ---------------------------------------------------------------------------
# Setup wizard
# ---------------------------------------------------------------------------
class SetupWizard:
"""
Interactive onboarding wizard: provider selection → API keys → config file.
@ -194,9 +219,14 @@ class SetupWizard:
port, master_key = SetupWizard._proxy_settings()
config_path = Path(os.getcwd()) / "litellm_config.yaml"
config_path.write_text(
SetupWizard._build_config(providers, env_vars, port, master_key)
)
try:
config_path.write_text(
SetupWizard._build_config(providers, env_vars, port, master_key)
)
except OSError as exc:
print(f"\n {bold(_CROSS + ' Could not write config:')} {exc}")
print(" Try running from a directory you have write access to.\n")
return
SetupWizard._print_success(config_path, port, master_key)
SetupWizard._offer_start(config_path, port, master_key)
@ -213,7 +243,7 @@ class SetupWizard:
print(orange(LITELLM_ASCII.rstrip("\n")))
print(f" {orange('Welcome')} to {bold('LiteLLM')} {grey('v' + version)}")
print()
print(DIVIDER)
print(_divider())
print()
# ── provider selector ───────────────────────────────────────────────────
@ -310,7 +340,7 @@ class SetupWizard:
print()
while True:
raw = input(f" {blue('❯')} Provider(s): ").strip()
raw = _styled_input(f" {blue('❯')} Provider(s): ")
if not raw:
print(grey(" Please select at least one provider."))
continue
@ -330,10 +360,11 @@ class SetupWizard:
def _collect_keys(providers: List[Dict]) -> Dict[str, str]:
env_vars: Dict[str, str] = {}
print()
print(DIVIDER)
print(_divider())
print()
print(f" {bold('Enter your API keys')}")
print(grey(" Keys are stored only in the generated config file."))
print(grey(" Tip: add litellm_config.yaml to .gitignore to avoid committing secrets."))
print()
for p in providers:
@ -350,16 +381,21 @@ class SetupWizard:
for extra_key, extra_hint in zip(
p.get("extra_keys", []), p.get("extra_hints", [])
):
val = input(f" {blue('❯')} {extra_key} {grey(extra_hint)}: ").strip()
val = _styled_input(f" {blue('❯')} {extra_key} {grey(extra_hint)}: ")
if val:
env_vars[extra_key] = val
if p.get("needs_api_base"):
api_base = input(
api_base = _styled_input(
f" {blue('❯')} Azure endpoint URL {grey(p.get('api_base_hint', ''))}: "
).strip()
)
if api_base:
env_vars[f"_LITELLM_AZURE_API_BASE_{p['id'].upper()}"] = api_base
deployment = _styled_input(
f" {blue('❯')} Azure deployment name {grey('(e.g. my-gpt4o)')}: "
)
if deployment:
env_vars[f"_LITELLM_AZURE_DEPLOYMENT_{p['id'].upper()}"] = deployment
SetupWizard._validate_and_report(p, key)
@ -370,17 +406,17 @@ class SetupWizard:
"""Prompt for a provider's API key, with skip option. Returns the key or ''."""
hint = grey(provider.get("key_hint", ""))
while True:
key = input(f" {blue('❯')} {bold(provider['name'])} API key {hint}: ").strip()
key = _styled_input(f" {blue('❯')} {bold(provider['name'])} API key {hint}: ")
if key:
return key
print(grey(" Key is required. Leave blank to skip this provider."))
if input(grey(" Skip? (y/N): ")).strip().lower() == "y":
if _styled_input(grey(" Skip? (y/N): ")).lower() == "y":
return ""
@staticmethod
def _validate_and_report(provider: Dict, api_key: str) -> None:
"""
Validate credentials using litellm.utils.check_valid_key.
Validate credentials using litellm.utils.check_valid_key and print result.
Offers a re-entry loop on failure.
"""
test_model: Optional[str] = provider.get("test_model")
@ -388,20 +424,20 @@ class SetupWizard:
return # Azure / Bedrock / Ollama — skip
while True:
print(f" {grey('Testing credentials…')}", end="", flush=True)
print(f" {grey('Testing connection to ' + provider['name'] + '...')}", flush=True)
valid = check_valid_key(model=test_model, api_key=api_key)
if valid:
print(f"\r {green(_CHECK + ' ' + provider['name'])} credentials valid ")
print(f" {green(_CHECK)} {bold(provider['name'])} connected successfully")
return
print(f"\r {_c(_BOLD, _CROSS)} {bold(provider['name'])} {grey('invalid API key')}")
if input(f" {blue('❯')} Re-enter key? {grey('(y/N)')}: ").strip().lower() != "y":
print(f" {_CROSS} {bold(provider['name'])} {grey('— invalid API key')}")
if _styled_input(f" {blue('❯')} Re-enter key? {grey('(y/N)')}: ").lower() != "y":
return
hint = grey(provider.get("key_hint", ""))
new_key = input(
new_key = _styled_input(
f" {blue('❯')} {bold(provider['name'])} API key {hint}: "
).strip()
)
if not new_key:
return
api_key = new_key
@ -411,13 +447,20 @@ class SetupWizard:
@staticmethod
def _proxy_settings() -> "tuple[int, str]":
print()
print(DIVIDER)
print(_divider())
print()
print(f" {bold('Proxy settings')}")
print()
port_raw = input(f" {blue('❯')} Port {grey('[4000]')}: ").strip()
port = int(port_raw) if port_raw.isdigit() else 4000
key_raw = input(f" {blue('❯')} Master key {grey('[auto-generate]')}: ").strip()
port = 4000
while True:
port_raw = _styled_input(f" {blue('❯')} Port {grey('[4000]')}: ")
if not port_raw:
break
if port_raw.isdigit() and 1 <= int(port_raw) <= 65535:
port = int(port_raw)
break
print(grey(" Enter a valid port number (1–65535)."))
key_raw = _styled_input(f" {blue('❯')} Master key {grey('[auto-generate]')}: ")
master_key = key_raw if key_raw else f"sk-{secrets.token_urlsafe(32)}"
return port, master_key
@ -430,34 +473,49 @@ class SetupWizard:
port: int,
master_key: str,
) -> str:
env_copy = dict(env_vars) # work on a copy — do not mutate caller's dict
lines = ["model_list:"]
for p in providers:
models = p["models"] if p["models"] else (["azure/gpt-4o"] if p["id"] == "azure" else [])
# Only emit models for providers that actually have credentials
has_creds = p["env_key"] is None or p["env_key"] in env_copy
if not has_creds:
continue
if p["id"] == "azure":
deployment = env_copy.pop(
f"_LITELLM_AZURE_DEPLOYMENT_{p['id'].upper()}", "gpt-4o"
)
models = [f"azure/{deployment}"]
else:
models = p["models"]
for model in models:
display = model.split("/")[-1] if "/" in model else model
raw_display = model.split("/")[-1] if "/" in model else model
# Qualify azure display names to avoid collision with OpenAI model names
display = f"azure-{raw_display}" if p["id"] == "azure" else raw_display
lines += [
f" - model_name: {display}",
f" litellm_params:",
" litellm_params:",
f" model: {model}",
]
if p["env_key"] and p["env_key"] in env_vars:
if p["env_key"] and p["env_key"] in env_copy:
lines.append(f" api_key: os.environ/{p['env_key']}")
if p.get("api_base"):
lines.append(f" api_base: {p['api_base']}")
elif p.get("needs_api_base"):
azure_key = f"_LITELLM_AZURE_API_BASE_{p['id'].upper()}"
if azure_key in env_vars:
lines.append(f" api_base: {env_vars.pop(azure_key)}")
azure_base_key = f"_LITELLM_AZURE_API_BASE_{p['id'].upper()}"
if azure_base_key in env_copy:
lines.append(f" api_base: {env_copy.pop(azure_base_key)}")
if p.get("api_version"):
lines.append(f" api_version: {p['api_version']}")
lines += ["", "general_settings:", f" master_key: {master_key}", ""]
real_vars = {k: v for k, v in env_vars.items() if not k.startswith("_LITELLM_")}
real_vars = {k: v for k, v in env_copy.items() if not k.startswith("_LITELLM_")}
if real_vars:
lines.append("environment_variables:")
for k, v in real_vars.items():
lines.append(f' {k}: "{v}"')
lines.append(f' {k}: "{_yaml_escape(v)}"')
lines.append("")
return "\n".join(lines)
@ -467,7 +525,7 @@ class SetupWizard:
@staticmethod
def _print_success(config_path: Path, port: int, master_key: str) -> None:
print()
print(DIVIDER)
print(_divider())
print()
print(f" {green(_CHECK + ' Config saved')} → {bold(str(config_path))}")
print()
@ -480,12 +538,14 @@ class SetupWizard:
print(f" export OPENAI_BASE_URL=http://localhost:{port}")
print(f" export OPENAI_API_KEY={master_key}")
print()
print(DIVIDER)
print(_divider())
print()
@staticmethod
def _offer_start(config_path: Path, port: int, master_key: str) -> None:
start = input(f" {blue('❯')} Start the proxy now? {grey('(Y/n)')}: ").strip().lower()
start = _styled_input(
f" {blue('❯')} Start the proxy now? {grey('(Y/n)')}: "
).lower()
if start not in ("", "y", "yes"):
print()
print(f" Run {bold(f'litellm --config {config_path}')} whenever you're ready.")
@ -495,7 +555,7 @@ class SetupWizard:
return
print()
print(DIVIDER)
print(_divider())
print()
print(f" {bold('Proxy is starting on')} http://localhost:{port}")
print()
@ -512,21 +572,25 @@ class SetupWizard:
print()
print(f" http://localhost:{port}/ui {grey('(login with your master key)')}")
print()
print(DIVIDER)
print(_divider())
print()
print(f" {green(_CHECK)} Starting… {grey('(Ctrl+C to stop)')}")
print()
scripts_dir = sysconfig.get_path("scripts")
litellm_bin = os.path.join(scripts_dir, "litellm")
os.execlp(litellm_bin, litellm_bin, "--config", str(config_path), "--port", str(port)) # noqa: S606
litellm_bin = os.path.join(scripts_dir or "", "litellm")
try:
os.execlp(litellm_bin, litellm_bin, "--config", str(config_path), "--port", str(port)) # noqa: S606
except OSError as exc:
print(f"\n {bold(_CROSS + ' Could not start proxy:')} {exc}")
print(f" Run manually: litellm --config {config_path} --port {port}\n")
# ---------------------------------------------------------------------------
# Public entrypoint
# ---------------------------------------------------------------------------
def run_setup_wizard() -> Optional[str]:
def run_setup_wizard() -> None:
"""Run the interactive setup wizard. Called by `litellm --setup`."""
SetupWizard.run()
return None

19
scripts/install.sh
View file

@ -1,14 +1,16 @@
#!/usr/bin/env bash
# LiteLLM Installer
# Usage: curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh
set -euo pipefail
#
# NOTE: set -e without pipefail for POSIX sh compatibility (dash on Ubuntu/Debian
# ignores the shebang when invoked as `sh` and does not support `pipefail`).
set -eu
MIN_PYTHON_MAJOR=3
MIN_PYTHON_MINOR=9
# NOTE: set to "litellm[proxy]" before merging to main (once --setup ships in a PyPI release).
# On this branch we install from git so the installer gets the --setup flag.
LITELLM_PACKAGE="litellm[proxy] @ git+https://github.com/BerriAI/litellm.git@worktree-dynamic-tickling-journal"
# NOTE: before merging, this must stay as "litellm[proxy]" to install from PyPI.
LITELLM_PACKAGE="litellm[proxy]"
# ── colours ────────────────────────────────────────────────────────────────
if [ -t 1 ]; then
@ -79,10 +81,9 @@ fi
# ── install ────────────────────────────────────────────────────────────────
echo ""
header "Installing litellm[proxy]…"
printf " \033[38;2;177;185;249m ℹ This clones from GitHub and may take 2–3 minutes on first install.\033[0m\n"
echo ""
"$PYTHON_BIN" -m pip install --upgrade --force-reinstall "${LITELLM_PACKAGE}" \
"$PYTHON_BIN" -m pip install --upgrade "${LITELLM_PACKAGE}" \
|| die "pip install failed. Try manually: $PYTHON_BIN -m pip install '${LITELLM_PACKAGE}'"
# ── find the litellm binary installed by pip for this Python ───────────────
@ -116,7 +117,11 @@ fi
# ── launch setup wizard ────────────────────────────────────────────────────
echo ""
printf " ${BOLD}Run the interactive setup wizard?${RESET} ${GREY}(Y/n)${RESET}: "
read -r answer </dev/tty
# /dev/tty may be unavailable in Docker/CI — default to yes if it can't be read
answer=""
if [ -r /dev/tty ]; then
read -r answer </dev/tty || answer=""
fi
if [ -z "$answer" ] || [ "$answer" = "y" ] || [ "$answer" = "Y" ]; then
echo ""

159
tests/test_litellm/test_setup_wizard.py Normal file
View file

@ -0,0 +1,159 @@
"""Unit tests for litellm.setup_wizard — pure functions only, no network calls."""
from litellm.setup_wizard import SetupWizard, _yaml_escape
# ---------------------------------------------------------------------------
# _yaml_escape
# ---------------------------------------------------------------------------
def test_yaml_escape_plain():
assert _yaml_escape("sk-abc123") == "sk-abc123"
def test_yaml_escape_double_quote():
assert _yaml_escape('sk-ab"cd') == 'sk-ab\\"cd'
def test_yaml_escape_backslash():
assert _yaml_escape("sk-ab\\cd") == "sk-ab\\\\cd"
def test_yaml_escape_combined():
assert _yaml_escape('ab\\"cd') == 'ab\\\\\\"cd'
# ---------------------------------------------------------------------------
# SetupWizard._build_config
# ---------------------------------------------------------------------------
_OPENAI = {
"id": "openai",
"name": "OpenAI",
"env_key": "OPENAI_API_KEY",
"models": ["gpt-4o", "gpt-4o-mini"],
"test_model": "gpt-4o-mini",
}
_ANTHROPIC = {
"id": "anthropic",
"name": "Anthropic",
"env_key": "ANTHROPIC_API_KEY",
"models": ["claude-opus-4-6"],
"test_model": "claude-haiku-4-5-20251001",
}
_AZURE = {
"id": "azure",
"name": "Azure OpenAI",
"env_key": "AZURE_API_KEY",
"models": [],
"test_model": None,
"needs_api_base": True,
"api_base_hint": "https://<resource>.openai.azure.com/",
"api_version": "2024-07-01-preview",
}
_OLLAMA = {
"id": "ollama",
"name": "Ollama",
"env_key": None,
"models": ["ollama/llama3.2"],
"test_model": None,
"api_base": "http://localhost:11434",
}
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
def test_build_config_skipped_provider_omitted():
"""Provider with no key in env_vars should not appear in model_list."""
config = SetupWizard._build_config(
[_OPENAI, _ANTHROPIC],
{"ANTHROPIC_API_KEY": "sk-ant-test"}, # OpenAI key missing
4000,
"sk-master",
)
assert "gpt-4o" not in config
assert "claude-opus-4-6" in config
def test_build_config_env_vars_written_escaped():
"""API keys with special chars should be YAML-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_does_not_mutate_env_vars():
"""_build_config must not modify the caller's env_vars dict."""
env_vars = {
"AZURE_API_KEY": "az-key",
"_LITELLM_AZURE_API_BASE_AZURE": "https://my.azure.com",
"_LITELLM_AZURE_DEPLOYMENT_AZURE": "my-deployment",
}
original_keys = set(env_vars.keys())
SetupWizard._build_config([_AZURE], env_vars, 4000, "sk-master")
assert set(env_vars.keys()) == original_keys
def test_build_config_azure_uses_deployment_name():
env_vars = {
"AZURE_API_KEY": "az-key",
"_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")
assert "model: azure/my-gpt4o" in config
assert "model_name: azure-my-gpt4o" in config
def test_build_config_no_display_name_collision_openai_and_azure():
"""OpenAI gpt-4o and azure gpt-4o should get distinct model_name values."""
env_vars = {
"OPENAI_API_KEY": "sk-openai",
"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
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")
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_internal_sentinel_keys_excluded():
"""_LITELLM_ prefixed sentinel keys must not appear in environment_variables."""
env_vars = {
"OPENAI_API_KEY": "sk-real",
"_LITELLM_AZURE_API_BASE_AZURE": "https://x.azure.com",
}
config = SetupWizard._build_config([_OPENAI], env_vars, 4000, "sk-master")
assert "_LITELLM_" not in config