Files
MoFin/venv/lib/python3.12/site-packages/litellm/setup_wizard.py
T
知微 fa45d8aa5f fix: 小果地址统一node122(兼容LAN+EasyTier)
- health_checklist.json: 192.168.1.122→node122
- ocr_client.py: docstring IP→node122
- docs/market-data-requirements.md: IP→node122
- 所有API调用通过ProxyHandler({})绕过系统代理
  Privoxy对node122:18003返回500,直连正常
2026-06-30 02:56:35 +08:00

676 lines
24 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# ruff: noqa: T201
# flake8: noqa: T201
"""
LiteLLM Interactive Setup Wizard
Guides users through selecting LLM providers, entering API keys,
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
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
# ---------------------------------------------------------------------------
# Provider definitions
# ---------------------------------------------------------------------------
# Each entry describes one provider card shown in the wizard.
# `env_key` — primary env var name (None = no key needed, e.g. Ollama)
# `test_model` — model passed to check_valid_key for credential validation
# (None = skip validation, e.g. Azure needs a deployment name)
# `models` — default models written into the generated config
# ---------------------------------------------------------------------------
PROVIDERS: List[Dict] = [
{
"id": "openai",
"name": "OpenAI",
"description": "GPT-4o, GPT-4o-mini, o3-mini",
"env_key": "OPENAI_API_KEY",
"key_hint": "sk-...",
"test_model": "gpt-4o-mini",
"models": ["gpt-4o", "gpt-4o-mini"],
},
{
"id": "anthropic",
"name": "Anthropic",
"description": "Claude Fable 5, Opus 4.8, Opus 4.7, Opus 4.6, Sonnet 4.6, Haiku 4.5",
"env_key": "ANTHROPIC_API_KEY",
"key_hint": "sk-ant-...",
"test_model": "claude-haiku-4-5-20251001",
"models": [
"claude-fable-5",
"claude-opus-4-8",
"claude-opus-4-7",
"claude-opus-4-6",
"claude-sonnet-4-6",
"claude-haiku-4-5-20251001",
],
},
{
"id": "gemini",
"name": "Google Gemini",
"description": "Gemini 2.0 Flash, Gemini 2.5 Pro",
"env_key": "GEMINI_API_KEY",
"key_hint": "AIza...",
"test_model": "gemini/gemini-2.0-flash",
"models": ["gemini/gemini-2.0-flash", "gemini/gemini-2.5-pro"],
},
{
"id": "azure",
"name": "Azure OpenAI",
"description": "GPT-4o via Azure",
"env_key": "AZURE_AI_API_KEY",
"key_hint": "your-azure-key",
"test_model": None, # needs deployment name — skip validation
"models": [],
"needs_api_base": True,
"api_base_hint": "https://<resource>.openai.azure.com/",
"api_version": "2024-07-01-preview",
},
{
"id": "bedrock",
"name": "AWS Bedrock",
"description": "Claude 3.5, Llama 3 via AWS",
"env_key": "AWS_ACCESS_KEY_ID",
"key_hint": "AKIA...",
"test_model": None, # multi-key auth — skip validation
"models": ["bedrock/anthropic.claude-haiku-4-5-20251001-v1:0"],
"extra_keys": ["AWS_SECRET_ACCESS_KEY", "AWS_REGION_NAME"],
"extra_hints": ["your-secret-key", "us-east-1"],
},
{
"id": "ollama",
"name": "Ollama",
"description": "Local models (llama3.2, mistral, etc.)",
"env_key": None,
"key_hint": None,
"test_model": None, # local — no remote validation
"models": ["ollama/llama3.2", "ollama/mistral"],
"api_base": "http://localhost:11434",
},
]
# ---------------------------------------------------------------------------
# ANSI colour helpers
# ---------------------------------------------------------------------------
_ANSI_RE = re.compile(r"\033\[[^m]*m")
_ORANGE = "\033[38;2;215;119;87m"
_DIM = "\033[2m"
_BOLD = "\033[1m"
_GREEN = "\033[38;2;78;186;101m"
_BLUE = "\033[38;2;177;185;249m"
_GREY = "\033[38;2;153;153;153m"
_RESET = "\033[0m"
_CHECK = ""
_CROSS = ""
_CURSOR_HIDE = "\033[?25l"
_CURSOR_SHOW = "\033[?25h"
_MOVE_UP = "\033[{}A"
def _supports_color() -> bool:
return sys.stdout.isatty() and os.environ.get("NO_COLOR") is None
def _c(code: str, text: str) -> str:
return f"{code}{text}{_RESET}" if _supports_color() else text
def orange(t: str) -> str:
return _c(_ORANGE, t)
def bold(t: str) -> str:
return _c(_BOLD, t)
def green(t: str) -> str:
return _c(_GREEN, t)
def blue(t: str) -> str:
return _c(_BLUE, t)
def grey(t: str) -> str:
return _c(_GREY, t)
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('"', '\\"')
.replace("\n", "\\n")
.replace("\r", "\\r")
.replace("\t", "\\t")
)
# ---------------------------------------------------------------------------
# Layout constants
# ---------------------------------------------------------------------------
LITELLM_ASCII = r"""
██╗ ██╗████████╗███████╗██╗ ██╗ ███╗ ███╗
██║ ██║╚══██╔══╝██╔════╝██║ ██║ ████╗ ████║
██║ ██║ ██║ █████╗ ██║ ██║ ██╔████╔██║
██║ ██║ ██║ ██╔══╝ ██║ ██║ ██║╚██╔╝██║
███████╗██║ ██║ ███████╗███████╗███████╗██║ ╚═╝ ██║
╚══════╝╚═╝ ╚═╝ ╚══════╝╚══════╝╚══════╝╚═╝ ╚═╝
"""
# ---------------------------------------------------------------------------
# Setup wizard
# ---------------------------------------------------------------------------
class SetupWizard:
"""
Interactive onboarding wizard: provider selection → API keys → config file.
All methods are static — the class is purely a namespace with clear
single-responsibility sections. Entry point: SetupWizard.run().
"""
# ── entry point ─────────────────────────────────────────────────────────
@staticmethod
def run() -> None:
try:
SetupWizard._wizard()
except (KeyboardInterrupt, EOFError):
print(f"\n\n {grey('Setup cancelled.')}\n")
# ── wizard steps ────────────────────────────────────────────────────────
@staticmethod
def _wizard() -> None:
SetupWizard._print_welcome()
print(f" {bold('Lets get started.')}")
print()
providers = SetupWizard._select_providers()
env_vars = SetupWizard._collect_keys(providers)
port, master_key = SetupWizard._proxy_settings()
config_path = Path(os.getcwd()) / "litellm_config.yaml"
try:
config_path.write_text(
SetupWizard._build_config(providers, env_vars, 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)
# ── welcome ─────────────────────────────────────────────────────────────
@staticmethod
def _print_welcome() -> None:
try:
version = importlib.metadata.version("litellm")
except Exception:
version = "unknown"
print()
print(orange(LITELLM_ASCII.rstrip("\n")))
print(f" {orange('Welcome')} to {bold('LiteLLM')} {grey('v' + version)}")
print()
print(_divider())
print()
# ── provider selector ───────────────────────────────────────────────────
@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:
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)
try:
tty.setraw(fd)
ch = tty_fh.read(1)
if ch == b"\x1b":
ch2 = tty_fh.read(1)
if ch2 == b"[":
ch3 = tty_fh.read(1)
return "\x1b[" + ch3.decode("utf-8", errors="replace")
return "\x1b" + ch2.decode("utf-8", errors="replace")
return ch.decode("utf-8", errors="replace")
finally:
termios.tcsetattr(fd, termios.TCSADRAIN, old)
@staticmethod
def _render_selector(cursor: int, selected: Set[int], first_render: bool) -> int:
"""Draw or redraw the provider list. Returns the number of lines printed."""
lines = [
f"\n {bold('Add your first model')}\n",
grey(" ↑↓ to navigate · Space to select · Enter to confirm") + "\n",
"\n",
]
for i, p in enumerate(PROVIDERS):
arrow = blue("") if i == cursor else " "
bullet = green("") if i in selected else grey("")
name_str = bold(p["name"]) if i == cursor else p["name"]
lines.append(f" {arrow} {bullet} {name_str} {grey(p['description'])}\n")
lines.append("\n")
content = "".join(lines)
if not first_render and _supports_color():
sys.stdout.write(_MOVE_UP.format(content.count("\n")))
sys.stdout.write(content)
sys.stdout.flush()
return content.count("\n")
@staticmethod
def _select_interactive() -> List[Dict]:
cursor = 0
selected: set[int] = set()
if _supports_color():
sys.stdout.write(_CURSOR_HIDE)
sys.stdout.flush()
try:
SetupWizard._render_selector(cursor, selected, first_render=True)
while True:
key = SetupWizard._read_key()
dirty = False
if key == "\x1b[A":
cursor = (cursor - 1) % len(PROVIDERS)
dirty = True
elif key == "\x1b[B":
cursor = (cursor + 1) % len(PROVIDERS)
dirty = True
elif key == " ":
selected.symmetric_difference_update({cursor})
dirty = True
elif key in ("\r", "\n"):
if not selected:
selected.add(cursor)
break
elif key in ("\x03", "\x04"):
raise KeyboardInterrupt
if dirty:
SetupWizard._render_selector(cursor, selected, first_render=False)
finally:
if _supports_color():
sys.stdout.write(_CURSOR_SHOW)
sys.stdout.flush()
return [PROVIDERS[i] for i in sorted(selected)]
@staticmethod
def _select_fallback() -> List[Dict]:
"""Number-based fallback when raw terminal input is unavailable."""
print()
print(f" {bold('Add your first model')}")
print(
grey(
" Enter numbers separated by commas (e.g. 1,2). Press Enter to confirm."
)
)
print()
for i, p in enumerate(PROVIDERS, 1):
print(f" {grey(str(i) + '.')} {bold(p['name'])} {grey(p['description'])}")
print()
while True:
raw = _styled_input(f" {blue('')} Provider(s): ")
if not raw:
print(grey(" Please select at least one provider."))
continue
try:
nums = [
int(x.strip())
for x in raw.replace(" ", ",").split(",")
if x.strip()
]
valid = sorted({n for n in nums if 1 <= n <= len(PROVIDERS)})
if not valid:
print(grey(f" Enter numbers between 1 and {len(PROVIDERS)}."))
continue
return [PROVIDERS[i - 1] for i in valid]
except ValueError:
print(grey(" Enter numbers separated by commas, e.g. 1,3"))
# ── key collection ───────────────────────────────────────────────────────
@staticmethod
def _collect_keys(providers: List[Dict]) -> Dict[str, str]:
env_vars: Dict[str, str] = {}
print()
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:
if p["env_key"] is None:
print(
f" {green(p['name'])}: {grey('no key needed (uses local Ollama)')}"
)
continue
key = SetupWizard._prompt_key(p)
if not key:
continue
for extra_key, extra_hint in zip(
p.get("extra_keys", []), p.get("extra_hints", [])
):
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 = _styled_input(
f" {blue('')} Azure endpoint URL {grey(p.get('api_base_hint', ''))}: "
)
if api_base:
env_vars[f"_LITELLM_AZURE_AI_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
)
# 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
@staticmethod
def _prompt_key(provider: Dict) -> str:
"""Prompt for a provider's API key, with skip option. Returns the key or ''."""
hint = grey(provider.get("key_hint", ""))
while True:
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 _styled_input(grey(" Skip? (y/N): ")).lower() == "y":
return ""
@staticmethod
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. Returns the final (possibly re-entered) key.
"""
test_model: Optional[str] = provider.get("test_model")
if not test_model:
return api_key # Azure / Bedrock / Ollama — skip validation
while 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" {green(_CHECK)} {bold(provider['name'])} connected successfully"
)
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 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 api_key
api_key = new_key
# ── proxy settings ───────────────────────────────────────────────────────
@staticmethod
def _proxy_settings() -> "tuple[int, str]":
print()
print(_divider())
print()
print(f" {bold('Proxy settings')}")
print()
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 (165535)."))
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
# ── config generation ────────────────────────────────────────────────────
@staticmethod
def _build_config(
providers: List[Dict],
env_vars: Dict[str, str],
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:
# 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()}", ""
)
if not deployment:
continue # skip Azure entirely if no deployment name was provided
models = [f"azure/{deployment}"]
else:
models = p["models"]
for model in models:
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}",
" litellm_params:",
f" model: {model}",
]
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: "{_yaml_escape(str(p["api_base"]))}"'
)
elif p.get("needs_api_base"):
azure_base_key = f"_LITELLM_AZURE_AI_API_BASE_{p['id'].upper()}"
if azure_base_key in env_copy:
lines.append(
f' api_base: "{_yaml_escape(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: "{_yaml_escape(master_key)}"',
"",
]
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}: "{_yaml_escape(v)}"')
lines.append("")
return "\n".join(lines)
# ── success + launch ─────────────────────────────────────────────────────
@staticmethod
def _print_success(config_path: Path, port: int, master_key: str) -> None:
print()
print(_divider())
print()
print(f" {green(_CHECK + ' Config saved')}{bold(str(config_path))}")
print()
print(f" {bold('To start your proxy:')}")
print()
print(f" {grey('$')} litellm --config {config_path} --port {port}")
print()
print(f" {bold('Then set your client:')}")
print()
print(f" export OPENAI_BASE_URL=http://localhost:{port}")
print(f" export OPENAI_API_KEY={master_key}")
print()
print(_divider())
print()
@staticmethod
def _offer_start(config_path: Path, port: int, master_key: str) -> None:
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."
)
print()
print(
grey(f" Quick test once running: curl http://localhost:{port}/health")
)
print()
return
print()
print(_divider())
print()
print(f" {bold('Proxy is starting on')} http://localhost:{port}")
print()
print(grey(" Your proxy is OpenAI-compatible. Point any OpenAI SDK at it:"))
print()
print(f" export OPENAI_BASE_URL=http://localhost:{port}")
print(f" export OPENAI_API_KEY={master_key}")
print()
print(grey(" Quick test (in another terminal):"))
print()
print(f" curl http://localhost:{port}/health")
print()
print(grey(" Dashboard:"))
print()
print(f" http://localhost:{port}/ui {grey('(login with your master key)')}")
print()
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 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() -> None:
"""Run the interactive setup wizard. Called by `litellm --setup`."""
SetupWizard.run()