From 60d3d057047b311b49e63521f6124b72e6711ab3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 14 Jul 2026 11:10:41 -0700 Subject: [PATCH] feat(cli): add lite autoroute to QA complexity-based auto-routing against a real proxy Lets a customer try litellm's complexity_router against models they already have on their existing, unmodified production proxy, with no config.yaml edits and no new infra. lite autoroute configure discovers accessible models via /model_group/info and walks through tier assignment (plus optional LLM classifier / semantic matching / adaptive selection); every referenced model becomes its own litellm_proxy/ deployment forwarding back to the real proxy with the real key, so every actual call, routed completions, classifier calls, embedding calls, still lands on their real proxy. lite autoroute up launches that generated config as an ephemeral local proxy, patches ~/.claude/settings.json to point Claude Code at it, and streams routing decisions live; Ctrl-C/SIGTERM (or lite autoroute down after an unclean exit) restores everything. Also adds lite model-groups list (a thin CLI wrapper over the existing ModelGroupsManagementClient), and generalizes up.py's settings-backup/restore helpers to take explicit paths so this feature can reuse them instead of duplicating the logic. Depends on litellm_lite_up_down (#33231) for that generalization. --- litellm/proxy/client/cli/README.md | 55 +++++ .../client/cli/commands/autoroute/__init__.py | 0 .../client/cli/commands/autoroute/commands.py | 157 ++++++++++++ .../client/cli/commands/autoroute/config.py | 223 ++++++++++++++++++ .../client/cli/commands/autoroute/process.py | 158 +++++++++++++ .../client/cli/commands/autoroute/settings.py | 35 +++ .../client/cli/commands/autoroute/wizard.py | 106 +++++++++ .../proxy/client/cli/commands/model_groups.py | 54 +++++ litellm/proxy/client/cli/commands/up.py | 34 +-- litellm/proxy/client/cli/main.py | 6 + .../proxy/client/cli/autoroute/__init__.py | 0 .../client/cli/autoroute/test_commands.py | 175 ++++++++++++++ .../proxy/client/cli/autoroute/test_config.py | 167 +++++++++++++ .../client/cli/autoroute/test_process.py | 113 +++++++++ .../client/cli/autoroute/test_settings.py | 36 +++ .../proxy/client/cli/autoroute/test_wizard.py | 188 +++++++++++++++ .../client/cli/test_model_groups_commands.py | 104 ++++++++ 17 files changed, 1597 insertions(+), 14 deletions(-) create mode 100644 litellm/proxy/client/cli/commands/autoroute/__init__.py create mode 100644 litellm/proxy/client/cli/commands/autoroute/commands.py create mode 100644 litellm/proxy/client/cli/commands/autoroute/config.py create mode 100644 litellm/proxy/client/cli/commands/autoroute/process.py create mode 100644 litellm/proxy/client/cli/commands/autoroute/settings.py create mode 100644 litellm/proxy/client/cli/commands/autoroute/wizard.py create mode 100644 litellm/proxy/client/cli/commands/model_groups.py create mode 100644 tests/test_litellm/proxy/client/cli/autoroute/__init__.py create mode 100644 tests/test_litellm/proxy/client/cli/autoroute/test_commands.py create mode 100644 tests/test_litellm/proxy/client/cli/autoroute/test_config.py create mode 100644 tests/test_litellm/proxy/client/cli/autoroute/test_process.py create mode 100644 tests/test_litellm/proxy/client/cli/autoroute/test_settings.py create mode 100644 tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py create mode 100644 tests/test_litellm/proxy/client/cli/test_model_groups_commands.py diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index e8b024dc31d..ecd5f3d48e2 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -489,6 +489,61 @@ This is a one-time file patch and restore, not a live traffic interceptor. A Cla Cursor is not supported: it has no equivalent file-based config to hot-patch this way, since its model routing lives in its own app storage and is configured through its GUI. +### QA Complexity-Based Auto-Routing Against Your Real Proxy + +`lite autoroute` lets you try LiteLLM's complexity-based auto-routing -- picking a cheaper or more expensive model depending on how complex a prompt looks -- against models your key already has access to on your real, running proxy, without editing that proxy's `config.yaml` and without any real request ever bypassing it. It builds a second, throwaway proxy locally that forwards every request back to your real proxy, and points Claude Code at that local proxy for the duration of the session. + +#### List Your Accessible Model Groups + +```bash +lite model-groups list [--format table|json] +``` + +Lists the model groups your key can reach on the proxy, via `/model_group/info`, along with each group's mode (`chat`, `embedding`, etc.) and per-token pricing. This is also what `lite autoroute configure` uses internally to discover what it can offer you. + +#### Configure the Auto-Router + +```bash +lite autoroute configure +``` + +An interactive wizard. It runs the same model-group discovery as above, splits the results into chat-capable and embedding-capable pools, and asks you to assign a model from the chat pool to each of the four complexity tiers -- SIMPLE, MEDIUM, COMPLEX, REASONING. From there it optionally offers: classifying prompt complexity with an LLM (again picked from your discovered pool) instead of the free built-in heuristic scorer, semantic keyword matching for tier assignment (needs an embedding model from the pool), and adaptive (bandit-based) selection layered on top of tiering. + +The wizard writes the result to `~/.litellm/autorouter/config.yaml` with `0600` permissions, since the file embeds your real proxy API key. Every model referenced anywhere in that config -- tier targets, the classifier model, the embedding model -- becomes its own `litellm_proxy/` deployment whose `api_base` and `api_key` point back at your real proxy. That is the trick that keeps your real proxy's config untouched: every actual network call this generates, whether it is the routed completion, an LLM-classifier call, or an embedding call, forwards transparently through your real, already-running proxy with your real key. + +You must run `configure` at least once before `up`; running `up` first fails with a clear error telling you to configure first. + +#### Launch the Ephemeral Auto-Router Proxy + +```bash +lite autoroute up +``` + +Starts a local, throwaway litellm proxy on a random free port, running the config `configure` generated, with a freshly-minted random API key baked in for this session only (your real proxy key never leaves the generated config -- it only appears there, forwarding to your real proxy). It waits for the ephemeral proxy to report healthy, then patches `~/.claude/settings.json` the same way `lite up` does, except with a static `ANTHROPIC_AUTH_TOKEN` env var instead of an `apiKeyHelper`, since this key is short-lived and self-issued rather than something needing SSO refresh. Any `claude` session started afterward, from any terminal, routes through the ephemeral proxy. + +`lite autoroute up` runs in the foreground and streams the ephemeral proxy's own log file into your terminal, so you can watch its routing decisions -- which tier and model got picked for each request -- as you use Claude Code normally. Press Ctrl-C (or send SIGTERM) to stop it; this kills the child proxy process and restores your original Claude Code settings, in that order. + +#### Recover From an Unclean Shutdown + +```bash +lite autoroute down +``` + +If the `lite autoroute up` process dies uncleanly -- `kill -9`, a crash -- rather than being stopped with Ctrl-C, `down` is the manual recovery path: it kills any leftover ephemeral proxy process found via a recorded pid file and restores Claude Code's settings from whatever backup is on disk. + +#### Example + +```bash +lite autoroute configure +lite autoroute up +# use Claude Code as normal in another terminal; routing decisions stream live +lite autoroute down # only needed if `up` was killed uncleanly instead of Ctrl-C'd +``` + +#### Caveats + +Adaptive mode's learned state does not persist across `lite autoroute up` sessions -- there is no local database, so every session starts adaptive selection cold. A Claude Code session already running before `up` started, or still running when it stops, keeps whatever settings it loaded at its own startup; like `lite up`, this is a one-time file patch and restore, not a live traffic interceptor. Only Claude Code is supported, for the same reason as `lite up`: no other supported agent (for example Cursor) has an equivalent hot-patchable config file. + ## Environment Variables The CLI respects the following environment variables: diff --git a/litellm/proxy/client/cli/commands/autoroute/__init__.py b/litellm/proxy/client/cli/commands/autoroute/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/client/cli/commands/autoroute/commands.py b/litellm/proxy/client/cli/commands/autoroute/commands.py new file mode 100644 index 00000000000..aca2f395615 --- /dev/null +++ b/litellm/proxy/client/cli/commands/autoroute/commands.py @@ -0,0 +1,157 @@ +import atexit +import json +import secrets +import signal +import threading +from types import FrameType +from typing import Dict, Optional + +import click +import yaml +from pydantic import JsonValue, TypeAdapter + +from ..up import CLAUDE_SETTINGS_PATH +from ..up import BackupRecord as ClaudeBackupRecord +from ..up import load_json_or_empty, restore_claude_settings, write_backup +from .process import ( + AUTOROUTE_DIR, + CONFIG_PATH, + LOG_PATH, + PidRecord, + ProcessLaunchError, + allocate_free_port, + clear_pid_record, + is_running, + launch_proxy, + poll_liveliness, + read_pid_record, + stream_log, + terminate, + write_pid_record, +) +from .settings import merge_claude_settings_static_token +from .wizard import run_configure_wizard + +AUTOROUTE_BACKUP_PATH = AUTOROUTE_DIR / "claude_settings_backup.json" + +_GENERATED_CONFIG_ADAPTER = TypeAdapter(Dict[str, JsonValue]) + + +def _mint_and_embed_master_key() -> str: + """Generate a fresh key for this session and write it into the generated config.yaml. + + Must go under general_settings, not litellm_settings -- the proxy server only ever + reads general_settings.master_key (proxy_server.py:4530) to authenticate requests. A + key placed under litellm_settings is silently ignored, leaving the ephemeral proxy with + no real auth: any request reaches it regardless of the token Claude Code sends. + """ + master_key = secrets.token_urlsafe(32) + with open(CONFIG_PATH, "r") as f: + generated = _GENERATED_CONFIG_ADAPTER.validate_python(yaml.safe_load(f)) + general_settings = generated.get("general_settings") + updated_settings: Dict[str, JsonValue] = { + **(general_settings if isinstance(general_settings, dict) else {}), + "master_key": master_key, + } + updated: Dict[str, JsonValue] = {**generated, "general_settings": updated_settings} + with open(CONFIG_PATH, "w") as f: + yaml.safe_dump(updated, f, sort_keys=False) + CONFIG_PATH.chmod(0o600) + return master_key + + +@click.group(name="autoroute") +def autoroute_group() -> None: + """QA complexity-based auto-routing against models your key can already use""" + + +@autoroute_group.command("configure") +@click.pass_context +def configure(ctx: click.Context) -> None: + """Discover accessible models and generate an ephemeral auto-router config""" + run_configure_wizard(ctx) + + +@autoroute_group.command("up") +def up() -> None: + """Launch the ephemeral auto-router proxy and route Claude Code through it""" + if not CONFIG_PATH.exists(): + raise click.ClickException("No config found. Run `lite autoroute configure` first.") + + existing_pid = read_pid_record() + if existing_pid is not None and is_running(existing_pid.pid): + raise click.ClickException( + "An ephemeral proxy is already running (lite autoroute up looks already active). " + "Run `lite autoroute down` first." + ) + + master_key = _mint_and_embed_master_key() + port = allocate_free_port() + base_url = f"http://127.0.0.1:{port}" + process = launch_proxy(CONFIG_PATH, port, LOG_PATH) + write_pid_record(PidRecord(pid=process.pid, port=port, config_path=str(CONFIG_PATH), log_path=str(LOG_PATH))) + + try: + poll_liveliness(base_url, LOG_PATH, process) + except ProcessLaunchError as e: + clear_pid_record() + raise click.ClickException(str(e)) + + original_existed = CLAUDE_SETTINGS_PATH.exists() + original_settings = load_json_or_empty(CLAUDE_SETTINGS_PATH) + write_backup( + ClaudeBackupRecord(existed=original_existed, content=original_settings if original_existed else None), + AUTOROUTE_BACKUP_PATH, + ) + merged = merge_claude_settings_static_token(original_settings, base_url, master_key) + CLAUDE_SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True) + with open(CLAUDE_SETTINGS_PATH, "w") as f: + json.dump(merged, f, indent=2) + + click.echo(f"litellm: ephemeral auto-router proxy up at {base_url} (pid {process.pid})") + click.echo("Claude Code sessions started now will route through it. Press Ctrl-C to stop and restore.") + + stop_event = threading.Event() + restored = threading.Lock() + + def _teardown() -> None: + if not restored.acquire(blocking=False): + return + terminate(process.pid) + clear_pid_record() + restore_claude_settings(CLAUDE_SETTINGS_PATH, AUTOROUTE_BACKUP_PATH) + click.echo("\nStopped ephemeral proxy and restored Claude Code settings.") + + def _handle_signal(_signum: int, _frame: Optional[FrameType]) -> None: + stop_event.set() + + signal.signal(signal.SIGINT, _handle_signal) + signal.signal(signal.SIGTERM, _handle_signal) + atexit.register(_teardown) + + log_thread = threading.Thread(target=stream_log, args=(LOG_PATH, stop_event), daemon=True) + log_thread.start() + + stop_event.wait() + _teardown() + + +@autoroute_group.command("down") +def down() -> None: + """Restore Claude Code settings and stop a leftover ephemeral proxy, if any""" + record: Optional[PidRecord] = read_pid_record() + if record is not None and is_running(record.pid): + terminate(record.pid) + click.echo(f"Stopped leftover ephemeral proxy (pid {record.pid}).") + clear_pid_record() + + restored = restore_claude_settings(CLAUDE_SETTINGS_PATH, AUTOROUTE_BACKUP_PATH) + if restored is None: + click.echo("Nothing to restore.") + elif restored.existed: + click.echo(f"Restored {CLAUDE_SETTINGS_PATH} to its original contents.") + else: + click.echo(f"Removed {CLAUDE_SETTINGS_PATH} (it did not exist before `lite autoroute up`).") + + +__all__ = ["autoroute_group"] diff --git a/litellm/proxy/client/cli/commands/autoroute/config.py b/litellm/proxy/client/cli/commands/autoroute/config.py new file mode 100644 index 00000000000..85b5223524b --- /dev/null +++ b/litellm/proxy/client/cli/commands/autoroute/config.py @@ -0,0 +1,223 @@ +from typing import Dict, FrozenSet, List, Literal, Tuple, Union + +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter + +TIER_NAMES: Tuple[str, ...] = ("SIMPLE", "MEDIUM", "COMPLEX", "REASONING") + + +class ConfigGenerationError(Exception): + """Raised when an AutorouteConfig references a model the discovery step didn't find.""" + + +class DiscoveredModel(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str + mode: str = "chat" + input_cost_per_token: float | None = None + output_cost_per_token: float | None = None + + +class _RawModelGroup(BaseModel): + model_config = ConfigDict(extra="ignore") + + model_group: str + # Optional: some real deployments return an explicit `"mode": null` for models that + # were registered without a mode (seen for embedding models like voyage-4-large). + # ModelGroupInfo's own "chat" default (litellm/types/router.py) only applies when the + # key is missing entirely, not when it's present as null, so this must tolerate None. + mode: str | None = "chat" + input_cost_per_token: float | None = None + output_cost_per_token: float | None = None + + +_RAW_MODEL_GROUPS_ADAPTER = TypeAdapter(List[_RawModelGroup]) + + +def parse_discovered_models(raw: List[JsonValue]) -> Tuple[DiscoveredModel, ...]: + """Validate a raw `/model_group/info` response into typed models.""" + parsed = _RAW_MODEL_GROUPS_ADAPTER.validate_python(raw) + return tuple( + DiscoveredModel( + name=group.model_group, + # A null mode means the server genuinely doesn't know what this model does; + # "unknown" (rather than guessing "chat") keeps it out of both chat_models() + # and embedding_models() instead of risking a wrong-mode deployment. + mode=group.mode or "unknown", + input_cost_per_token=group.input_cost_per_token, + output_cost_per_token=group.output_cost_per_token, + ) + for group in parsed + ) + + +def chat_models(models: Tuple[DiscoveredModel, ...]) -> Tuple[DiscoveredModel, ...]: + return tuple(m for m in models if m.mode == "chat") + + +def embedding_models(models: Tuple[DiscoveredModel, ...]) -> Tuple[DiscoveredModel, ...]: + return tuple(m for m in models if m.mode == "embedding") + + +class HeuristicClassifier(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["heuristic"] = "heuristic" + + +class LLMClassifier(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["llm"] = "llm" + model: str + timeout_ms: int = 3000 + + +ClassifierChoice = Union[HeuristicClassifier, LLMClassifier] + + +class NoSemanticMatching(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["none"] = "none" + + +class SemanticMatching(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["semantic"] = "semantic" + embedding_model: str + match_threshold: float = 0.5 + + +SemanticMatchingChoice = Union[NoSemanticMatching, SemanticMatching] + +# Satisfies complexity_router's "semantic matching requires non-empty keyword_tier_rules" +# invariant with a sane starting point; the generated config.yaml can be hand-edited afterward. +_DEFAULT_KEYWORD_TIER_RULES: Tuple[Dict[str, JsonValue], ...] = ( + {"keywords": ["hi", "hello", "thanks"], "tier": "SIMPLE"}, + {"keywords": ["explain", "how does"], "tier": "MEDIUM"}, + {"keywords": ["refactor", "implement", "debug"], "tier": "COMPLEX"}, + {"keywords": ["step by step", "think through", "prove"], "tier": "REASONING"}, +) + + +class AutorouteConfig(BaseModel): + model_config = ConfigDict(frozen=True) + + base_url: str + api_key: str + tiers: Dict[str, str] + default_model: str + classifier: ClassifierChoice = Field(default_factory=HeuristicClassifier) + semantic_matching: SemanticMatchingChoice = Field(default_factory=NoSemanticMatching) + adaptive: bool = False + + +def validate_config(config: AutorouteConfig, discovered: Tuple[DiscoveredModel, ...]) -> None: + """Raise ConfigGenerationError if config references a model discovery didn't return.""" + chat_names: FrozenSet[str] = frozenset(m.name for m in chat_models(discovered)) + embedding_names: FrozenSet[str] = frozenset(m.name for m in embedding_models(discovered)) + + for tier, model in config.tiers.items(): + if model not in chat_names: + raise ConfigGenerationError(f"Tier {tier} references unknown chat model '{model}'") + + if config.default_model not in chat_names: + raise ConfigGenerationError(f"default_model '{config.default_model}' is not a known chat model") + + if isinstance(config.classifier, LLMClassifier) and config.classifier.model not in chat_names: + raise ConfigGenerationError(f"classifier model '{config.classifier.model}' is not a known chat model") + + if ( + isinstance(config.semantic_matching, SemanticMatching) + and config.semantic_matching.embedding_model not in embedding_names + ): + raise ConfigGenerationError( + f"embedding model '{config.semantic_matching.embedding_model}' is not a known embedding model" + ) + + +def _litellm_proxy_deployment(name: str, base_url: str, api_key: str) -> Dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": { + "model": f"litellm_proxy/{name}", + "api_base": base_url, + "api_key": api_key, + }, + } + + +def build_generated_model_list(config: AutorouteConfig) -> List[JsonValue]: + """Build the model_list for the ephemeral proxy's config.yaml. + + Every real model referenced anywhere (tier targets, classifier, embedding) is deduplicated + to exactly one `litellm_proxy/` deployment forwarding to the customer's real proxy, + plus one `auto_router/complexity_router` deployment tying the tiers together. + """ + referenced_names = {*config.tiers.values(), config.default_model} + if isinstance(config.classifier, LLMClassifier): + referenced_names.add(config.classifier.model) + if isinstance(config.semantic_matching, SemanticMatching): + referenced_names.add(config.semantic_matching.embedding_model) + + proxy_deployments = [ + _litellm_proxy_deployment(name, config.base_url, config.api_key) for name in sorted(referenced_names) + ] + + complexity_router_config: Dict[str, JsonValue] = { + "tiers": dict(config.tiers), + "default_model": config.default_model, + } + if isinstance(config.classifier, LLMClassifier): + complexity_router_config["classifier_type"] = "llm" + complexity_router_config["classifier_llm_config"] = { + "model": config.classifier.model, + "timeout_ms": config.classifier.timeout_ms, + } + if isinstance(config.semantic_matching, SemanticMatching): + complexity_router_config["semantic_keyword_matching"] = True + complexity_router_config["embedding_model"] = config.semantic_matching.embedding_model + complexity_router_config["match_threshold"] = config.semantic_matching.match_threshold + complexity_router_config["keyword_tier_rules"] = list(_DEFAULT_KEYWORD_TIER_RULES) + if config.adaptive: + complexity_router_config["adaptive"] = True + + auto_router_deployment: Dict[str, JsonValue] = { + "model_name": "autorouter", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": complexity_router_config, + }, + } + return [*proxy_deployments, auto_router_deployment] + + +def build_generated_proxy_config(config: AutorouteConfig, master_key: str) -> Dict[str, JsonValue]: + """Full config.yaml content for the ephemeral proxy, including its own auth key. + + master_key must live under general_settings, not litellm_settings -- the proxy server + only ever reads general_settings.master_key (proxy_server.py:4530) to authenticate + requests; a key placed under litellm_settings is silently ignored, leaving the proxy + with no real auth at all. + """ + return { + "model_list": build_generated_model_list(config), + "general_settings": {"master_key": master_key}, + } + + +__all__ = [ + "TIER_NAMES", + "ConfigGenerationError", + "DiscoveredModel", + "parse_discovered_models", + "chat_models", + "embedding_models", + "HeuristicClassifier", + "LLMClassifier", + "ClassifierChoice", + "NoSemanticMatching", + "SemanticMatching", + "SemanticMatchingChoice", + "AutorouteConfig", + "validate_config", + "build_generated_model_list", +] diff --git a/litellm/proxy/client/cli/commands/autoroute/process.py b/litellm/proxy/client/cli/commands/autoroute/process.py new file mode 100644 index 00000000000..ce146e95eec --- /dev/null +++ b/litellm/proxy/client/cli/commands/autoroute/process.py @@ -0,0 +1,158 @@ +import contextlib +import json +import os +import signal +import socket +import subprocess +import sys +import threading +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +import click +import requests +from pydantic import TypeAdapter + +AUTOROUTE_DIR = Path.home() / ".litellm" / "autorouter" +CONFIG_PATH = AUTOROUTE_DIR / "config.yaml" +LOG_PATH = AUTOROUTE_DIR / "proxy.log" +PID_RECORD_PATH = AUTOROUTE_DIR / "proxy.pid.json" + + +class ProcessLaunchError(Exception): + """Raised when the ephemeral proxy subprocess fails to come up healthy.""" + + +@dataclass(frozen=True, slots=True) +class PidRecord: + pid: int + port: int + config_path: str + log_path: str + + +_PID_RECORD_ADAPTER = TypeAdapter(PidRecord) + + +def allocate_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def launch_proxy(config_path: Path, port: int, log_path: Path) -> "subprocess.Popen[bytes]": + log_path.parent.mkdir(parents=True, exist_ok=True) + log_file = open(log_path, "w") + return subprocess.Popen( + [sys.executable, "-m", "litellm.proxy.proxy_cli", "--config", str(config_path), "--port", str(port)], + stdout=log_file, + stderr=subprocess.STDOUT, + ) + + +def _tail(log_path: Path, lines: int = 40) -> str: + if not log_path.exists(): + return "(no log output captured)" + return "\n".join(log_path.read_text(errors="replace").splitlines()[-lines:]) + + +def poll_liveliness(base_url: str, log_path: Path, process: "subprocess.Popen[bytes]", timeout: float = 30.0) -> None: + """Poll /health/liveliness until it responds, the process dies, or timeout elapses.""" + deadline = time.monotonic() + timeout + url = base_url.rstrip("/") + "/health/liveliness" + while time.monotonic() < deadline: + if process.poll() is not None: + raise ProcessLaunchError( + f"Ephemeral proxy exited early (code {process.returncode}). Last log lines:\n{_tail(log_path)}" + ) + with contextlib.suppress(requests.RequestException): + if requests.get(url, timeout=2).status_code == 200: + return + time.sleep(0.5) + raise ProcessLaunchError( + f"Ephemeral proxy never became healthy within {timeout}s. Last log lines:\n{_tail(log_path)}" + ) + + +def write_pid_record(record: PidRecord, path: Optional[Path] = None) -> None: + resolved_path = path if path is not None else PID_RECORD_PATH + resolved_path.parent.mkdir(parents=True, exist_ok=True) + with open(resolved_path, "w") as f: + json.dump( + {"pid": record.pid, "port": record.port, "config_path": record.config_path, "log_path": record.log_path}, + f, + indent=2, + ) + + +def read_pid_record(path: Optional[Path] = None) -> Optional[PidRecord]: + resolved_path = path if path is not None else PID_RECORD_PATH + if not resolved_path.exists(): + return None + with open(resolved_path, "r") as f: + return _PID_RECORD_ADAPTER.validate_json(f.read()) + + +def clear_pid_record(path: Optional[Path] = None) -> None: + resolved_path = path if path is not None else PID_RECORD_PATH + resolved_path.unlink(missing_ok=True) + + +def is_running(pid: int) -> bool: + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def terminate(pid: int, grace_period: float = 5.0) -> None: + """Terminate a process by pid, escalating from SIGTERM to SIGKILL if needed.""" + if not is_running(pid): + return + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGTERM) + deadline = time.monotonic() + grace_period + while time.monotonic() < deadline and is_running(pid): + time.sleep(0.2) + if is_running(pid): + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + + +def stream_log(log_path: Path, stop_event: threading.Event) -> None: + """Print new lines appended to log_path until stop_event is set. Blocks the calling thread.""" + while not log_path.exists() and not stop_event.is_set(): + time.sleep(0.1) + if stop_event.is_set() or not log_path.exists(): + return + with open(log_path, "r") as f: + while not stop_event.is_set(): + line = f.readline() + if line: + click.echo(line, nl=False) + else: + time.sleep(0.2) + + +__all__ = [ + "AUTOROUTE_DIR", + "CONFIG_PATH", + "LOG_PATH", + "PID_RECORD_PATH", + "ProcessLaunchError", + "PidRecord", + "allocate_free_port", + "launch_proxy", + "poll_liveliness", + "write_pid_record", + "read_pid_record", + "clear_pid_record", + "is_running", + "terminate", + "stream_log", +] diff --git a/litellm/proxy/client/cli/commands/autoroute/settings.py b/litellm/proxy/client/cli/commands/autoroute/settings.py new file mode 100644 index 00000000000..0d7b12d7aaf --- /dev/null +++ b/litellm/proxy/client/cli/commands/autoroute/settings.py @@ -0,0 +1,35 @@ +from typing import Dict + +from pydantic import JsonValue + +ENV_KEY = "env" +API_KEY_HELPER_KEY = "apiKeyHelper" +ANTHROPIC_API_KEY_KEY = "ANTHROPIC_API_KEY" +ANTHROPIC_AUTH_TOKEN_KEY = "ANTHROPIC_AUTH_TOKEN" +ANTHROPIC_BASE_URL_KEY = "ANTHROPIC_BASE_URL" + + +def merge_claude_settings_static_token( + settings: Dict[str, JsonValue], base_url: str, auth_token: str +) -> Dict[str, JsonValue]: + """Return a new settings dict wired to a local ephemeral proxy with a static token. + + Unlike up.py's merge_claude_settings (which sets apiKeyHelper for a long-lived, real + remote proxy needing refreshable SSO tokens), this proxy is ephemeral and its key was just + minted for this session, so a plain env var is simpler and correct. Any existing + apiKeyHelper is cleared so it can't fight with the static token. + """ + raw_env = settings.get(ENV_KEY, {}) + base_env = raw_env if isinstance(raw_env, dict) else {} + env: Dict[str, JsonValue] = { + **base_env, + ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), + ANTHROPIC_AUTH_TOKEN_KEY: auth_token, + } + env.pop(ANTHROPIC_API_KEY_KEY, None) + merged: Dict[str, JsonValue] = {**settings, ENV_KEY: env} + merged.pop(API_KEY_HELPER_KEY, None) + return merged + + +__all__ = ["merge_claude_settings_static_token"] diff --git a/litellm/proxy/client/cli/commands/autoroute/wizard.py b/litellm/proxy/client/cli/commands/autoroute/wizard.py new file mode 100644 index 00000000000..e2d1b85dc6f --- /dev/null +++ b/litellm/proxy/client/cli/commands/autoroute/wizard.py @@ -0,0 +1,106 @@ +from pathlib import Path +from typing import Tuple + +import click +import yaml +from rich.console import Console +from rich.table import Table + +from .... import Client +from .config import ( + TIER_NAMES, + AutorouteConfig, + ConfigGenerationError, + DiscoveredModel, + HeuristicClassifier, + LLMClassifier, + NoSemanticMatching, + SemanticMatching, + build_generated_model_list, + chat_models, + embedding_models, + parse_discovered_models, + validate_config, +) +from .process import CONFIG_PATH + + +def _render_and_prompt_for_model(models: Tuple[DiscoveredModel, ...], prompt_label: str) -> str: + console = Console() + table = Table(title=f"Pick a model for {prompt_label}") + table.add_column("Index", style="cyan", no_wrap=True) + table.add_column("Model", style="magenta") + for i, model in enumerate(models): + table.add_row(str(i + 1), model.name) + console.print(table) + + while True: + choice = click.prompt(f"\nSelect a model for {prompt_label} by index", type=str).strip() + try: + index = int(choice) - 1 + except ValueError: + click.echo("Invalid input. Please enter a number.") + continue + if 0 <= index < len(models): + return models[index].name + click.echo(f"Invalid selection. Please enter a number between 1 and {len(models)}") + + +def run_configure_wizard(ctx: click.Context) -> Path: + """Discover the caller's accessible models, walk them through tier assignment, write config.""" + base_url = ctx.obj["base_url"] + api_key = ctx.obj["api_key"] + client = Client(base_url=base_url, api_key=api_key) + + raw_groups = client.model_groups.info() + assert isinstance(raw_groups, list) + discovered = parse_discovered_models(raw_groups) + chat_pool = chat_models(discovered) + embedding_pool = embedding_models(discovered) + + if not chat_pool: + raise click.ClickException("Your key has no chat-capable models available on this proxy.") + + click.echo("Assign a model to each complexity tier (from what your key can access):") + tiers = {tier: _render_and_prompt_for_model(chat_pool, tier) for tier in TIER_NAMES} + default_model = tiers["MEDIUM"] + + classifier = HeuristicClassifier() + if click.confirm("\nUse an LLM classifier instead of the free heuristic scorer?", default=False): + classifier_model = _render_and_prompt_for_model(chat_pool, "LLM classifier") + classifier = LLMClassifier(model=classifier_model) + + semantic_matching = NoSemanticMatching() + if embedding_pool and click.confirm("\nEnable semantic keyword matching?", default=False): + embedding_model = _render_and_prompt_for_model(embedding_pool, "semantic embeddings") + semantic_matching = SemanticMatching(embedding_model=embedding_model) + + adaptive = click.confirm("\nEnable adaptive (bandit) selection on top of tiering?", default=False) + + config = AutorouteConfig( + base_url=base_url, + api_key=api_key, + tiers=tiers, + default_model=default_model, + classifier=classifier, + semantic_matching=semantic_matching, + adaptive=adaptive, + ) + try: + validate_config(config, discovered) + except ConfigGenerationError as e: + raise click.ClickException(str(e)) + + model_list = build_generated_model_list(config) + CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True) + with open(CONFIG_PATH, "w") as f: + yaml.safe_dump({"model_list": model_list}, f, sort_keys=False) + CONFIG_PATH.chmod(0o600) + + click.echo(f"\nWrote {CONFIG_PATH}") + for tier, model in tiers.items(): + click.echo(f" {tier}: {model}") + return CONFIG_PATH + + +__all__ = ["run_configure_wizard"] diff --git a/litellm/proxy/client/cli/commands/model_groups.py b/litellm/proxy/client/cli/commands/model_groups.py new file mode 100644 index 00000000000..629a0b5aaf1 --- /dev/null +++ b/litellm/proxy/client/cli/commands/model_groups.py @@ -0,0 +1,54 @@ +from typing import Literal + +import click +import rich +import rich.table + +from ... import Client + + +def create_client(ctx: click.Context) -> Client: + return Client(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) + + +@click.group(name="model-groups") +def model_groups() -> None: + """Inspect model groups your key can access on the proxy""" + + +@model_groups.command("list") +@click.option( + "--format", + "output_format", + type=click.Choice(["table", "json"]), + default="table", + help="Output format (table or json)", +) +@click.pass_context +def list_model_groups(ctx: click.Context, output_format: Literal["table", "json"]) -> None: + """List model groups accessible to your key, with mode and pricing""" + client = create_client(ctx) + groups = client.model_groups.info() + assert isinstance(groups, list) + + if output_format == "json": + rich.print_json(data=groups) + return + + table = rich.table.Table(title="Accessible Model Groups") + table.add_column("Model", style="cyan") + table.add_column("Mode", style="green") + table.add_column("Input $/token", style="yellow") + table.add_column("Output $/token", style="yellow") + + for group in groups: + table.add_row( + str(group.get("model_group", "")), + str(group.get("mode", "chat")), + str(group.get("input_cost_per_token", "")), + str(group.get("output_cost_per_token", "")), + ) + rich.print(table) + + +__all__ = ["model_groups"] diff --git a/litellm/proxy/client/cli/commands/up.py b/litellm/proxy/client/cli/commands/up.py index fb9fb2b3fa4..b59a31147e3 100644 --- a/litellm/proxy/client/cli/commands/up.py +++ b/litellm/proxy/client/cli/commands/up.py @@ -68,34 +68,40 @@ def merge_claude_settings( return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper} -def write_backup(record: BackupRecord) -> None: - BACKUP_PATH.parent.mkdir(exist_ok=True) - with open(BACKUP_PATH, "w") as f: +def write_backup(record: BackupRecord, backup_path: Optional[Path] = None) -> None: + path = backup_path if backup_path is not None else BACKUP_PATH + path.parent.mkdir(exist_ok=True) + with open(path, "w") as f: json.dump({"existed": record.existed, "content": record.content}, f, indent=2) - os.chmod(BACKUP_PATH, 0o600) + os.chmod(path, 0o600) -def read_backup() -> Optional[BackupRecord]: - if not BACKUP_PATH.exists(): +def read_backup(backup_path: Optional[Path] = None) -> Optional[BackupRecord]: + path = backup_path if backup_path is not None else BACKUP_PATH + if not path.exists(): return None - with open(BACKUP_PATH, "r") as f: + with open(path, "r") as f: return _BACKUP_RECORD_ADAPTER.validate_json(f.read()) -def restore_claude_settings() -> Optional[BackupRecord]: - """Restore ~/.claude/settings.json from the backup, then delete the backup. +def restore_claude_settings( + settings_path: Optional[Path] = None, backup_path: Optional[Path] = None +) -> Optional[BackupRecord]: + """Restore settings_path from the backup at backup_path, then delete the backup. Returns the restored record, or None if there was nothing to restore. """ - record = read_backup() + resolved_settings_path = settings_path if settings_path is not None else CLAUDE_SETTINGS_PATH + resolved_backup_path = backup_path if backup_path is not None else BACKUP_PATH + record = read_backup(resolved_backup_path) if record is None: return None if record.existed and record.content is not None: - with open(CLAUDE_SETTINGS_PATH, "w") as f: + with open(resolved_settings_path, "w") as f: json.dump(record.content, f, indent=2) - elif CLAUDE_SETTINGS_PATH.exists(): - CLAUDE_SETTINGS_PATH.unlink() - BACKUP_PATH.unlink() + elif resolved_settings_path.exists(): + resolved_settings_path.unlink() + resolved_backup_path.unlink() return record diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index 322ce2b287a..e641956b2c5 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -9,11 +9,13 @@ from litellm.proxy.client.health import HealthManagementClient from .commands.agents import agent_commands from .commands.auth import auth_group, get_stored_api_key, login, logout, whoami +from .commands.autoroute.commands import autoroute_group from .commands.chat import chat from .commands.credentials import credentials from .commands.encryption import encryption from .commands.http import http from .commands.keys import keys +from .commands.model_groups import model_groups # local imports from .commands.models import models @@ -135,6 +137,10 @@ for agent_command in agent_commands(): # Add the up/down commands (route Claude Code through the local LiteLLM proxy) cli.add_command(up) cli.add_command(down) +# Add the model-groups command group (discover models your key can access) +cli.add_command(model_groups) +# Add the autoroute command group (QA auto-routing against your real proxy) +cli.add_command(autoroute_group, name="autoroute") if __name__ == "__main__": diff --git a/tests/test_litellm/proxy/client/cli/autoroute/__init__.py b/tests/test_litellm/proxy/client/cli/autoroute/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py new file mode 100644 index 00000000000..098659d5e8b --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py @@ -0,0 +1,175 @@ +import json +from typing import Optional + +import yaml +from click.testing import CliRunner + +from litellm.proxy.client.cli.commands.autoroute import commands as commands_module +from litellm.proxy.client.cli.commands.autoroute import process as process_module +from litellm.proxy.client.cli.commands.autoroute.commands import down, up +from litellm.proxy.client.cli.commands.autoroute.process import PidRecord, ProcessLaunchError, write_pid_record +from litellm.proxy.client.cli.commands.up import BackupRecord as ClaudeBackupRecord +from litellm.proxy.client.cli.commands.up import write_backup + + +class FakeProcess: + def __init__(self, pid: int): + self.pid = pid + self.returncode: Optional[int] = None + + def poll(self) -> Optional[int]: + return self.returncode + + +def _patch_paths(monkeypatch, tmp_path): + config_path = tmp_path / "config.yaml" + log_path = tmp_path / "proxy.log" + claude_settings_path = tmp_path / "claude_settings.json" + backup_path = tmp_path / "backup.json" + pid_record_path = tmp_path / "pid.json" + + monkeypatch.setattr(commands_module, "CONFIG_PATH", config_path) + monkeypatch.setattr(commands_module, "LOG_PATH", log_path) + monkeypatch.setattr(commands_module, "CLAUDE_SETTINGS_PATH", claude_settings_path) + monkeypatch.setattr(commands_module, "AUTOROUTE_BACKUP_PATH", backup_path) + monkeypatch.setattr(process_module, "PID_RECORD_PATH", pid_record_path) + + return config_path, log_path, claude_settings_path, backup_path, pid_record_path + + +def _silence_signal_handling(monkeypatch): + monkeypatch.setattr(commands_module.signal, "signal", lambda *a, **k: None) + monkeypatch.setattr(commands_module.atexit, "register", lambda *a, **k: None) + monkeypatch.setattr(commands_module, "stream_log", lambda *a, **k: None) + + +class TestUpCommand: + def setup_method(self): + self.runner = CliRunner() + + def test_refuses_when_never_configured(self, monkeypatch, tmp_path): + _patch_paths(monkeypatch, tmp_path) + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert "lite autoroute configure" in result.output + + def test_refuses_when_pid_record_exists_and_process_still_running(self, monkeypatch, tmp_path): + config_path, _log_path, _settings_path, _backup_path, pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + write_pid_record( + PidRecord(pid=123, port=4000, config_path=str(config_path), log_path="/tmp/proxy.log"), pid_record_path + ) + monkeypatch.setattr(commands_module, "is_running", lambda pid: True) + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert "already running" in result.output + assert "lite autoroute down" in result.output + assert config_path.read_text() == yaml.safe_dump({"model_list": []}) + + def test_happy_path_patches_settings_then_restores_everything_on_stop(self, monkeypatch, tmp_path): + config_path, log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + original_settings = {"theme": "dark"} + claude_settings_path.write_text(json.dumps(original_settings)) + _silence_signal_handling(monkeypatch) + + fake_process = FakeProcess(pid=99999) + terminate_calls = [] + monkeypatch.setattr(commands_module, "launch_proxy", lambda *a, **k: fake_process) + monkeypatch.setattr(commands_module, "poll_liveliness", lambda *a, **k: None) + monkeypatch.setattr(commands_module, "allocate_free_port", lambda: 54321) + monkeypatch.setattr(commands_module, "terminate", lambda pid, **k: terminate_calls.append(pid)) + monkeypatch.setattr(commands_module.secrets, "token_urlsafe", lambda n: "fixed-master-key") + + captured = {} + + def fake_wait(self, timeout=None): + captured["settings"] = json.loads(claude_settings_path.read_text()) + captured["backup_existed"] = backup_path.exists() + return True + + monkeypatch.setattr("threading.Event.wait", fake_wait) + + result = self.runner.invoke(up) + + assert result.exit_code == 0, result.output + assert captured["backup_existed"] is True + assert captured["settings"]["theme"] == "dark" + assert captured["settings"]["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:54321" + assert captured["settings"]["env"]["ANTHROPIC_AUTH_TOKEN"] == "fixed-master-key" + assert "apiKeyHelper" not in captured["settings"] + + assert terminate_calls == [99999] + assert not pid_record_path.exists() + assert not backup_path.exists() + assert json.loads(claude_settings_path.read_text()) == original_settings + + written_config = yaml.safe_load(config_path.read_text()) + assert written_config["general_settings"]["master_key"] == "fixed-master-key" + + def test_surfaces_clean_error_and_cleans_up_when_health_check_fails(self, monkeypatch, tmp_path): + config_path, _log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + original_settings = {"theme": "dark"} + claude_settings_path.write_text(json.dumps(original_settings)) + + fake_process = FakeProcess(pid=555) + + def _raise_launch_error(*args, **kwargs): + raise ProcessLaunchError("boom: proxy never became healthy") + + monkeypatch.setattr(commands_module, "launch_proxy", lambda *a, **k: fake_process) + monkeypatch.setattr(commands_module, "poll_liveliness", _raise_launch_error) + monkeypatch.setattr(commands_module, "allocate_free_port", lambda: 12345) + monkeypatch.setattr(commands_module.secrets, "token_urlsafe", lambda n: "fixed-master-key") + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert "boom" in result.output + assert not pid_record_path.exists() + assert not backup_path.exists() + assert json.loads(claude_settings_path.read_text()) == original_settings + + +class TestDownCommand: + def setup_method(self): + self.runner = CliRunner() + + def test_restores_settings_and_terminates_when_process_still_running(self, monkeypatch, tmp_path): + _config_path, _log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths( + monkeypatch, tmp_path + ) + original_settings = {"theme": "dark"} + write_backup(ClaudeBackupRecord(existed=True, content=original_settings), backup_path) + claude_settings_path.write_text(json.dumps({"env": {"ANTHROPIC_AUTH_TOKEN": "fixed-master-key"}})) + write_pid_record(PidRecord(pid=777, port=1234, config_path="c", log_path="l"), pid_record_path) + + terminate_calls = [] + monkeypatch.setattr(commands_module, "is_running", lambda pid: True) + monkeypatch.setattr(commands_module, "terminate", lambda pid, **k: terminate_calls.append(pid)) + + result = self.runner.invoke(down) + + assert result.exit_code == 0, result.output + assert "Stopped leftover ephemeral proxy" in result.output + assert "Restored" in result.output + assert terminate_calls == [777] + assert not pid_record_path.exists() + assert not backup_path.exists() + assert json.loads(claude_settings_path.read_text()) == original_settings + + def test_is_a_clean_no_op_when_nothing_is_running_and_no_backup_exists(self, monkeypatch, tmp_path): + _config_path, _log_path, claude_settings_path, _backup_path, _pid_record_path = _patch_paths( + monkeypatch, tmp_path + ) + + result = self.runner.invoke(down) + + assert result.exit_code == 0, result.output + assert "Nothing to restore." in result.output + assert not claude_settings_path.exists() diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_config.py b/tests/test_litellm/proxy/client/cli/autoroute/test_config.py new file mode 100644 index 00000000000..a48db0c79e4 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_config.py @@ -0,0 +1,167 @@ +from typing import Any, Dict, Tuple + +import pytest + +from litellm.proxy.client.cli.commands.autoroute.config import ( + AutorouteConfig, + ConfigGenerationError, + DiscoveredModel, + HeuristicClassifier, + LLMClassifier, + NoSemanticMatching, + SemanticMatching, + build_generated_model_list, + build_generated_proxy_config, + chat_models, + embedding_models, + parse_discovered_models, + validate_config, +) + +DISCOVERED: Tuple[DiscoveredModel, ...] = ( + DiscoveredModel(name="gpt-4o-mini", mode="chat"), + DiscoveredModel(name="gpt-4o", mode="chat"), + DiscoveredModel(name="o1", mode="chat"), + DiscoveredModel(name="text-embedding-3-small", mode="embedding"), +) + + +def _base_config(**overrides: Any) -> AutorouteConfig: + defaults: Dict[str, Any] = { + "base_url": "http://real-proxy.internal:4000", + "api_key": "sk-real-key", + "tiers": { + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "gpt-4o", + "COMPLEX": "gpt-4o", + "REASONING": "o1", + }, + "default_model": "gpt-4o", + } + defaults.update(overrides) + return AutorouteConfig(**defaults) + + +class TestParseDiscoveredModels: + def test_parses_valid_raw_list_into_typed_tuple(self): + raw = [ + { + "model_group": "gpt-4o", + "mode": "chat", + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + }, + {"model_group": "text-embedding-3-small", "mode": "embedding"}, + ] + result = parse_discovered_models(raw) + assert result == ( + DiscoveredModel(name="gpt-4o", mode="chat", input_cost_per_token=0.01, output_cost_per_token=0.02), + DiscoveredModel(name="text-embedding-3-small", mode="embedding"), + ) + + def test_ignores_unknown_extra_fields(self): + raw = [{"model_group": "gpt-4o", "mode": "chat", "totally_unknown_field": "whatever"}] + result = parse_discovered_models(raw) + assert result == (DiscoveredModel(name="gpt-4o", mode="chat"),) + + def test_missing_mode_defaults_to_chat(self): + raw = [{"model_group": "gpt-4o"}] + result = parse_discovered_models(raw) + assert result[0].mode == "chat" + + +class TestChatAndEmbeddingFiltering: + def test_filters_by_mode(self): + models = ( + DiscoveredModel(name="gpt-4o", mode="chat"), + DiscoveredModel(name="text-embedding-3-small", mode="embedding"), + DiscoveredModel(name="claude", mode="chat"), + ) + assert chat_models(models) == (models[0], models[2]) + assert embedding_models(models) == (models[1],) + + +class TestBuildGeneratedModelList: + def test_dedups_model_used_in_multiple_roles(self): + config = _base_config(classifier=LLMClassifier(model="gpt-4o")) + model_list = build_generated_model_list(config) + gpt4o_entries = [m for m in model_list if m["model_name"] == "gpt-4o"] + assert len(gpt4o_entries) == 1 + + def test_every_proxy_deployment_points_back_at_customer_proxy(self): + config = _base_config() + model_list = build_generated_model_list(config) + proxy_entries = [m for m in model_list if m["model_name"] != "autorouter"] + names = {m["model_name"] for m in proxy_entries} + assert names == {"gpt-4o-mini", "gpt-4o", "o1"} + for entry in proxy_entries: + assert entry["litellm_params"]["model"] == f"litellm_proxy/{entry['model_name']}" + assert entry["litellm_params"]["api_base"] == config.base_url + assert entry["litellm_params"]["api_key"] == config.api_key + + def test_complexity_router_config_reflects_llm_classifier(self): + config = _base_config(classifier=LLMClassifier(model="gpt-4o", timeout_ms=1234)) + autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") + router_config = autorouter["litellm_params"]["complexity_router_config"] + assert router_config["classifier_type"] == "llm" + assert router_config["classifier_llm_config"] == {"model": "gpt-4o", "timeout_ms": 1234} + assert "semantic_keyword_matching" not in router_config + assert "adaptive" not in router_config + + def test_complexity_router_config_reflects_semantic_matching(self): + config = _base_config( + semantic_matching=SemanticMatching(embedding_model="text-embedding-3-small", match_threshold=0.7) + ) + autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") + router_config = autorouter["litellm_params"]["complexity_router_config"] + assert router_config["semantic_keyword_matching"] is True + assert router_config["embedding_model"] == "text-embedding-3-small" + assert router_config["match_threshold"] == 0.7 + assert router_config["keyword_tier_rules"] + assert "classifier_type" not in router_config + + def test_complexity_router_config_reflects_adaptive(self): + config = _base_config(adaptive=True) + autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") + assert autorouter["litellm_params"]["complexity_router_config"]["adaptive"] is True + + def test_default_classifier_and_semantic_matching_add_no_extra_keys(self): + config = _base_config(classifier=HeuristicClassifier(), semantic_matching=NoSemanticMatching()) + autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") + router_config = autorouter["litellm_params"]["complexity_router_config"] + assert set(router_config.keys()) == {"tiers", "default_model"} + + +class TestBuildGeneratedProxyConfig: + def test_embeds_master_key_under_general_settings(self): + config = _base_config() + proxy_config = build_generated_proxy_config(config, "sk-master-123") + assert proxy_config["general_settings"] == {"master_key": "sk-master-123"} + assert proxy_config["model_list"] == build_generated_model_list(config) + + +class TestValidateConfig: + def test_passes_for_fully_valid_config(self): + validate_config(_base_config(), DISCOVERED) + + def test_raises_for_tier_referencing_unknown_model(self): + config = _base_config( + tiers={"SIMPLE": "unknown-model", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o", "REASONING": "o1"} + ) + with pytest.raises(ConfigGenerationError, match="unknown-model"): + validate_config(config, DISCOVERED) + + def test_raises_for_unknown_default_model(self): + config = _base_config(default_model="unknown-model") + with pytest.raises(ConfigGenerationError, match="unknown-model"): + validate_config(config, DISCOVERED) + + def test_raises_for_unknown_llm_classifier_model(self): + config = _base_config(classifier=LLMClassifier(model="unknown-model")) + with pytest.raises(ConfigGenerationError, match="unknown-model"): + validate_config(config, DISCOVERED) + + def test_raises_for_unknown_semantic_embedding_model(self): + config = _base_config(semantic_matching=SemanticMatching(embedding_model="unknown-embedding")) + with pytest.raises(ConfigGenerationError, match="unknown-embedding"): + validate_config(config, DISCOVERED) diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_process.py b/tests/test_litellm/proxy/client/cli/autoroute/test_process.py new file mode 100644 index 00000000000..bffd8b598f4 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_process.py @@ -0,0 +1,113 @@ +import os +import socket +from typing import Optional + +import pytest + +from litellm.proxy.client.cli.commands.autoroute import process as process_module +from litellm.proxy.client.cli.commands.autoroute.process import ( + PidRecord, + ProcessLaunchError, + allocate_free_port, + clear_pid_record, + is_running, + poll_liveliness, + read_pid_record, + write_pid_record, +) + + +class FakeProcess: + def __init__(self, returncode: Optional[int] = None): + self.returncode = returncode + + def poll(self) -> Optional[int]: + return self.returncode + + +class FakeResponse: + def __init__(self, status_code: int): + self.status_code = status_code + + +def test_allocate_free_port_returns_a_bindable_port(): + port = allocate_free_port() + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", port)) + + +class TestPidRecordRoundTrip: + def test_write_then_read_round_trips(self, tmp_path): + path = tmp_path / "pid.json" + record = PidRecord(pid=123, port=4000, config_path="/tmp/config.yaml", log_path="/tmp/proxy.log") + + write_pid_record(record, path) + + assert read_pid_record(path) == record + + def test_read_missing_file_returns_none(self, tmp_path): + assert read_pid_record(tmp_path / "missing.json") is None + + def test_clear_removes_an_existing_record(self, tmp_path): + path = tmp_path / "pid.json" + write_pid_record(PidRecord(pid=1, port=1, config_path="a", log_path="b"), path) + assert path.exists() + + clear_pid_record(path) + + assert not path.exists() + + def test_clear_missing_file_is_a_no_op(self, tmp_path): + clear_pid_record(tmp_path / "missing.json") + + def test_write_creates_parent_directories(self, tmp_path): + path = tmp_path / "nested" / "dir" / "pid.json" + + write_pid_record(PidRecord(pid=1, port=1, config_path="a", log_path="b"), path) + + assert path.exists() + + +class TestIsRunning: + def test_current_process_is_running(self): + assert is_running(os.getpid()) is True + + def test_huge_unlikely_pid_is_not_running(self): + assert is_running(2**30) is False + + def test_permission_error_from_kill_is_treated_as_running(self, monkeypatch): + def fake_kill(pid: int, sig: int) -> None: + raise PermissionError("not permitted to signal this pid") + + monkeypatch.setattr(process_module.os, "kill", fake_kill) + + assert is_running(999) is True + + +class TestPollLiveliness: + def test_succeeds_when_health_check_returns_200_quickly(self, monkeypatch, tmp_path): + monkeypatch.setattr(process_module.requests, "get", lambda url, timeout: FakeResponse(200)) + + poll_liveliness("http://127.0.0.1:4000", tmp_path / "proxy.log", FakeProcess(), timeout=5.0) + + def test_raises_with_log_tail_when_timeout_elapses(self, monkeypatch, tmp_path): + log_path = tmp_path / "proxy.log" + log_path.write_text("line one\nline two\nline three\n") + monkeypatch.setattr(process_module.requests, "get", lambda url, timeout: FakeResponse(500)) + monkeypatch.setattr(process_module.time, "sleep", lambda seconds: None) + + with pytest.raises(ProcessLaunchError) as exc_info: + poll_liveliness("http://127.0.0.1:4000", log_path, FakeProcess(), timeout=0.05) + + assert "never became healthy" in str(exc_info.value) + assert "line three" in str(exc_info.value) + + def test_raises_immediately_when_process_already_exited(self, tmp_path): + log_path = tmp_path / "proxy.log" + log_path.write_text("crash log line") + + with pytest.raises(ProcessLaunchError) as exc_info: + poll_liveliness("http://127.0.0.1:4000", log_path, FakeProcess(returncode=1), timeout=5.0) + + assert "exited early" in str(exc_info.value) + assert "crash log line" in str(exc_info.value) diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py new file mode 100644 index 00000000000..203d77b70c1 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py @@ -0,0 +1,36 @@ +from litellm.proxy.client.cli.commands.autoroute.settings import merge_claude_settings_static_token + + +def test_preserves_unrelated_top_level_keys(): + merged = merge_claude_settings_static_token({"theme": "dark"}, "http://127.0.0.1:4000", "token-abc") + assert merged["theme"] == "dark" + + +def test_preserves_unrelated_env_keys(): + settings = {"env": {"SOME_OTHER_VAR": "value"}} + merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") + assert merged["env"]["SOME_OTHER_VAR"] == "value" + + +def test_sets_base_url_and_auth_token(): + merged = merge_claude_settings_static_token({}, "http://127.0.0.1:4000/", "token-abc") + assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:4000" + assert merged["env"]["ANTHROPIC_AUTH_TOKEN"] == "token-abc" + + +def test_drops_stray_api_key(): + settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}} + merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") + assert "ANTHROPIC_API_KEY" not in merged["env"] + + +def test_removes_existing_api_key_helper(): + settings = {"apiKeyHelper": "/usr/local/bin/lite auth print-token"} + merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") + assert "apiKeyHelper" not in merged + + +def test_does_not_mutate_input(): + settings = {"env": {"FOO": "bar"}, "apiKeyHelper": "old-helper"} + merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") + assert settings == {"env": {"FOO": "bar"}, "apiKeyHelper": "old-helper"} diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py b/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py new file mode 100644 index 00000000000..81f54bf9b78 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py @@ -0,0 +1,188 @@ +from typing import Any, Dict, List, Tuple +from unittest.mock import patch + +import click +import yaml +from click.testing import CliRunner + +from litellm.proxy.client.cli.commands.autoroute import wizard as wizard_module +from litellm.proxy.client.cli.commands.autoroute.config import DiscoveredModel +from litellm.proxy.client.cli.commands.autoroute.wizard import ( + _render_and_prompt_for_model, + run_configure_wizard, +) + +CHAT_AND_EMBEDDING_GROUPS: List[Dict[str, Any]] = [ + {"model_group": "gpt-4o-mini", "mode": "chat", "input_cost_per_token": 0.01, "output_cost_per_token": 0.02}, + {"model_group": "gpt-4o", "mode": "chat", "input_cost_per_token": 0.01, "output_cost_per_token": 0.02}, + {"model_group": "claude-opus", "mode": "chat"}, + {"model_group": "o1", "mode": "chat"}, + {"model_group": "text-embedding-3-small", "mode": "embedding"}, +] + +CHAT_ONLY_GROUPS: List[Dict[str, Any]] = [ + {"model_group": "gpt-4o-mini", "mode": "chat"}, + {"model_group": "gpt-4o", "mode": "chat"}, + {"model_group": "claude-opus", "mode": "chat"}, + {"model_group": "o1", "mode": "chat"}, +] + +EMBEDDING_ONLY_GROUPS: List[Dict[str, Any]] = [ + {"model_group": "text-embedding-3-small", "mode": "embedding"}, +] + + +@click.command() +@click.pass_context +def _invoke_wizard(ctx: click.Context) -> None: + run_configure_wizard(ctx) + + +def _run(tmp_path, raw_groups: List[Dict[str, Any]], input_str: str): + config_path = tmp_path / "config.yaml" + runner = CliRunner() + with ( + patch.object(wizard_module, "Client") as mock_client_cls, + patch.object(wizard_module, "CONFIG_PATH", config_path), + ): + mock_client_cls.return_value.model_groups.info.return_value = raw_groups + result = runner.invoke( + _invoke_wizard, + obj={"base_url": "http://localhost:4000", "api_key": "sk-test"}, + input=input_str, + ) + return result, config_path + + +def _router_config(config_path) -> Dict[str, Any]: + written = yaml.safe_load(config_path.read_text()) + autorouter = next(m for m in written["model_list"] if m["model_name"] == "autorouter") + return autorouter["litellm_params"]["complexity_router_config"] + + +class TestRunConfigureWizardHappyPath: + def test_assigns_tiers_and_declines_everything(self, tmp_path): + result, config_path = _run( + tmp_path, + CHAT_AND_EMBEDDING_GROUPS, + input_str="1\n2\n3\n4\nn\nn\nn\n", + ) + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["tiers"] == { + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "gpt-4o", + "COMPLEX": "claude-opus", + "REASONING": "o1", + } + assert router_config["default_model"] == "gpt-4o" + assert "classifier_type" not in router_config + assert "classifier_llm_config" not in router_config + assert "semantic_keyword_matching" not in router_config + assert "adaptive" not in router_config + + def test_writes_config_file_with_restricted_permissions(self, tmp_path): + result, config_path = _run( + tmp_path, + CHAT_AND_EMBEDDING_GROUPS, + input_str="1\n2\n3\n4\nn\nn\nn\n", + ) + + assert result.exit_code == 0, result.output + assert config_path.exists() + assert oct(config_path.stat().st_mode)[-3:] == "600" + + def test_no_embedding_pool_skips_semantic_prompt_entirely(self, tmp_path): + result, config_path = _run( + tmp_path, + CHAT_ONLY_GROUPS, + input_str="1\n2\n3\n4\nn\nn\n", + ) + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert "semantic_keyword_matching" not in router_config + + +class TestRunConfigureWizardLLMClassifier: + def test_accepting_llm_classifier_records_chosen_model(self, tmp_path): + result, config_path = _run( + tmp_path, + CHAT_AND_EMBEDDING_GROUPS, + input_str="1\n2\n3\n4\ny\n2\nn\nn\n", + ) + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["classifier_type"] == "llm" + assert router_config["classifier_llm_config"]["model"] == "gpt-4o" + + +class TestRunConfigureWizardSemanticMatching: + def test_accepting_semantic_matching_records_embedding_model(self, tmp_path): + result, config_path = _run( + tmp_path, + CHAT_AND_EMBEDDING_GROUPS, + input_str="1\n2\n3\n4\nn\ny\n1\nn\n", + ) + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["semantic_keyword_matching"] is True + assert router_config["embedding_model"] == "text-embedding-3-small" + + +class TestRunConfigureWizardAdaptive: + def test_accepting_adaptive_sets_adaptive_flag(self, tmp_path): + result, config_path = _run( + tmp_path, + CHAT_AND_EMBEDDING_GROUPS, + input_str="1\n2\n3\n4\nn\nn\ny\n", + ) + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["adaptive"] is True + + +class TestRunConfigureWizardNoChatModels: + def test_fails_cleanly_without_prompting_when_no_chat_models(self, tmp_path): + result, config_path = _run(tmp_path, EMBEDDING_ONLY_GROUPS, input_str="") + + assert result.exit_code != 0 + assert "no chat-capable models" in result.output.lower() + assert not config_path.exists() + + +class TestRenderAndPromptForModel: + def _models(self) -> Tuple[DiscoveredModel, ...]: + return ( + DiscoveredModel(name="model-a"), + DiscoveredModel(name="model-b"), + ) + + def test_reprompts_on_non_numeric_input(self): + with patch("click.prompt", side_effect=["not-a-number", "2"]): + result = _render_and_prompt_for_model(self._models(), "test tier") + + assert result == "model-b" + + def test_reprompts_on_out_of_range_index(self): + with patch("click.prompt", side_effect=["5", "1"]): + result = _render_and_prompt_for_model(self._models(), "test tier") + + assert result == "model-a" + + def test_reprompts_on_zero_index(self): + with patch("click.prompt", side_effect=["0", "2"]): + result = _render_and_prompt_for_model(self._models(), "test tier") + + assert result == "model-b" + + def test_valid_first_answer_returns_immediately(self): + with patch("click.prompt", return_value="1") as mock_prompt: + result = _render_and_prompt_for_model(self._models(), "test tier") + + assert result == "model-a" + mock_prompt.assert_called_once() diff --git a/tests/test_litellm/proxy/client/cli/test_model_groups_commands.py b/tests/test_litellm/proxy/client/cli/test_model_groups_commands.py new file mode 100644 index 00000000000..d7b5b9dede4 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_model_groups_commands.py @@ -0,0 +1,104 @@ +import json +import os +from typing import Any, Dict, List +from unittest.mock import patch + +import pytest +from click.testing import CliRunner + +from litellm.proxy.client.cli import cli + +SAMPLE_MODEL_GROUPS: List[Dict[str, Any]] = [ + { + "model_group": "gpt-4o", + "mode": "chat", + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + }, + { + "model_group": "text-embedding-3-small", + "mode": "embedding", + "input_cost_per_token": 0.0001, + "output_cost_per_token": None, + }, +] + + +@pytest.fixture +def mock_client(): + with patch("litellm.proxy.client.cli.commands.model_groups.Client") as MockClient: + yield MockClient + + +@pytest.fixture +def cli_runner(): + return CliRunner() + + +@pytest.fixture(autouse=True) +def mock_env(): + with patch.dict( + os.environ, + { + "LITELLM_PROXY_URL": "http://localhost:4000", + "LITELLM_PROXY_API_KEY": "sk-test", + }, + ): + yield + + +def test_list_table_format_shows_model_names_and_modes(mock_client, cli_runner): + mock_client.return_value.model_groups.info.return_value = SAMPLE_MODEL_GROUPS + + result = cli_runner.invoke(cli, ["model-groups", "list"]) + + assert result.exit_code == 0, result.output + assert "gpt-4o" in result.output + assert "chat" in result.output + assert "text-embedding-3-small" in result.output + assert "embedding" in result.output + assert "0.01" in result.output + assert "0.02" in result.output + + mock_client.assert_called_once_with(base_url="http://localhost:4000", api_key="sk-test") + mock_client.return_value.model_groups.info.assert_called_once() + + +def test_list_table_format_defaults_missing_mode_to_chat(mock_client, cli_runner): + mock_client.return_value.model_groups.info.return_value = [{"model_group": "some-model"}] + + result = cli_runner.invoke(cli, ["model-groups", "list"]) + + assert result.exit_code == 0, result.output + assert "some-model" in result.output + assert "chat" in result.output + + +def test_list_json_format_round_trips_raw_data(mock_client, cli_runner): + mock_client.return_value.model_groups.info.return_value = SAMPLE_MODEL_GROUPS + + result = cli_runner.invoke(cli, ["model-groups", "list", "--format", "json"]) + + assert result.exit_code == 0, result.output + assert json.loads(result.output) == SAMPLE_MODEL_GROUPS + + +def test_list_with_custom_base_url_and_api_key(mock_client, cli_runner): + mock_client.return_value.model_groups.info.return_value = [] + + result = cli_runner.invoke( + cli, + ["--base-url", "http://custom.server:8000", "--api-key", "custom-key", "model-groups", "list"], + ) + + assert result.exit_code == 0, result.output + mock_client.assert_called_once_with(base_url="http://custom.server:8000", api_key="custom-key") + + +def test_list_error_handling(mock_client, cli_runner): + mock_client.return_value.model_groups.info.side_effect = Exception("API Error") + + result = cli_runner.invoke(cli, ["model-groups", "list"]) + + assert result.exit_code != 0 + assert "API Error" in str(result.exception)