chore: initial snapshot for gitea/github upload
This commit is contained in:
Your Name
@@ -0,0 +1,623 @@
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import webbrowser
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import click
|
||||
import requests
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
from litellm.constants import CLI_JWT_EXPIRATION_HOURS
|
||||
|
||||
|
||||
# Token storage utilities
|
||||
def get_token_file_path() -> str:
|
||||
"""Get the path to store the authentication token"""
|
||||
home_dir = Path.home()
|
||||
config_dir = home_dir / ".litellm"
|
||||
config_dir.mkdir(exist_ok=True)
|
||||
return str(config_dir / "token.json")
|
||||
|
||||
|
||||
def save_token(token_data: Dict[str, Any]) -> None:
|
||||
"""Save token data to file"""
|
||||
token_file = get_token_file_path()
|
||||
with open(token_file, "w") as f:
|
||||
json.dump(token_data, f, indent=2)
|
||||
# Set file permissions to be readable only by owner
|
||||
os.chmod(token_file, 0o600)
|
||||
|
||||
|
||||
def load_token() -> Optional[Dict[str, Any]]:
|
||||
"""Load token data from file"""
|
||||
token_file = get_token_file_path()
|
||||
if not os.path.exists(token_file):
|
||||
return None
|
||||
|
||||
try:
|
||||
with open(token_file, "r") as f:
|
||||
return json.load(f)
|
||||
except (json.JSONDecodeError, IOError):
|
||||
return None
|
||||
|
||||
|
||||
def clear_token() -> None:
|
||||
"""Clear stored token"""
|
||||
token_file = get_token_file_path()
|
||||
if os.path.exists(token_file):
|
||||
os.remove(token_file)
|
||||
|
||||
|
||||
def get_stored_api_key() -> Optional[str]:
|
||||
"""Get the stored API key from token file"""
|
||||
# Use the SDK-level utility
|
||||
from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_key
|
||||
|
||||
return get_litellm_gateway_api_key()
|
||||
|
||||
|
||||
# Team selection utilities
|
||||
def display_teams_table(teams: List[Dict[str, Any]]) -> None:
|
||||
"""Display teams in a formatted table"""
|
||||
console = Console()
|
||||
|
||||
if not teams:
|
||||
console.print("❌ No teams found for your user.")
|
||||
return
|
||||
|
||||
table = Table(title="Available Teams")
|
||||
table.add_column("Index", style="cyan", no_wrap=True)
|
||||
table.add_column("Team Alias", style="magenta")
|
||||
table.add_column("Team ID", style="green")
|
||||
table.add_column("Models", style="yellow")
|
||||
table.add_column("Max Budget", style="blue")
|
||||
|
||||
for i, team in enumerate(teams):
|
||||
team_alias = team.get("team_alias") or "N/A"
|
||||
team_id = team.get("team_id", "N/A")
|
||||
models = team.get("models", [])
|
||||
max_budget = team.get("max_budget")
|
||||
|
||||
# Format models list
|
||||
if models:
|
||||
if len(models) > 3:
|
||||
models_str = ", ".join(models[:3]) + f" (+{len(models) - 3} more)"
|
||||
else:
|
||||
models_str = ", ".join(models)
|
||||
else:
|
||||
models_str = "All models"
|
||||
|
||||
# Format budget
|
||||
budget_str = f"${max_budget}" if max_budget else "Unlimited"
|
||||
|
||||
table.add_row(str(i + 1), team_alias, team_id, models_str, budget_str)
|
||||
|
||||
console.print(table)
|
||||
|
||||
|
||||
def get_key_input():
|
||||
"""Get a single key input from the user (cross-platform)"""
|
||||
try:
|
||||
if sys.platform == "win32":
|
||||
import msvcrt
|
||||
|
||||
key = msvcrt.getch()
|
||||
if key == b"\xe0": # Arrow keys on Windows
|
||||
key = msvcrt.getch()
|
||||
if key == b"H": # Up arrow
|
||||
return "up"
|
||||
elif key == b"P": # Down arrow
|
||||
return "down"
|
||||
elif key == b"\r": # Enter key
|
||||
return "enter"
|
||||
elif key == b"\x1b": # Escape key
|
||||
return "escape"
|
||||
elif key == b"q":
|
||||
return "quit"
|
||||
return None
|
||||
else:
|
||||
import termios
|
||||
import tty
|
||||
|
||||
fd = sys.stdin.fileno()
|
||||
old_settings = termios.tcgetattr(fd)
|
||||
try:
|
||||
tty.setraw(sys.stdin.fileno())
|
||||
key = sys.stdin.read(1)
|
||||
|
||||
if key == "\x1b": # Escape sequence
|
||||
key += sys.stdin.read(2)
|
||||
if key == "\x1b[A": # Up arrow
|
||||
return "up"
|
||||
elif key == "\x1b[B": # Down arrow
|
||||
return "down"
|
||||
elif key == "\x1b": # Just escape
|
||||
return "escape"
|
||||
elif key == "\r" or key == "\n": # Enter key
|
||||
return "enter"
|
||||
elif key == "q":
|
||||
return "quit"
|
||||
return None
|
||||
finally:
|
||||
termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
|
||||
except ImportError:
|
||||
# Fallback to simple input if termios/msvcrt not available
|
||||
return None
|
||||
|
||||
|
||||
def display_interactive_team_selection(
|
||||
teams: List[Dict[str, Any]], selected_index: int = 0
|
||||
) -> None:
|
||||
"""Display teams with one highlighted for selection"""
|
||||
console = Console()
|
||||
|
||||
# Clear the screen using Rich's method
|
||||
console.clear()
|
||||
|
||||
console.print("🎯 Select a Team (Use ↑↓ arrows, Enter to select, 'q' to skip):\n")
|
||||
|
||||
for i, team in enumerate(teams):
|
||||
team_alias = team.get("team_alias") or "N/A"
|
||||
team_id = team.get("team_id", "N/A")
|
||||
models = team.get("models", [])
|
||||
max_budget = team.get("max_budget")
|
||||
|
||||
# Format models list
|
||||
if models:
|
||||
if len(models) > 3:
|
||||
models_str = ", ".join(models[:3]) + f" (+{len(models) - 3} more)"
|
||||
else:
|
||||
models_str = ", ".join(models)
|
||||
else:
|
||||
models_str = "All models"
|
||||
|
||||
# Format budget
|
||||
budget_str = f"${max_budget}" if max_budget else "Unlimited"
|
||||
|
||||
# Highlight the selected item
|
||||
if i == selected_index:
|
||||
console.print(f"➤ [bold cyan]{team_alias}[/bold cyan] ({team_id})")
|
||||
console.print(f" Models: [yellow]{models_str}[/yellow]")
|
||||
console.print(f" Budget: [blue]{budget_str}[/blue]\n")
|
||||
else:
|
||||
console.print(f" [dim]{team_alias}[/dim] ({team_id})")
|
||||
console.print(f" Models: [dim]{models_str}[/dim]")
|
||||
console.print(f" Budget: [dim]{budget_str}[/dim]\n")
|
||||
|
||||
|
||||
def prompt_team_selection(teams: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
|
||||
"""Interactive team selection with arrow keys"""
|
||||
if not teams:
|
||||
return None
|
||||
|
||||
selected_index = 0
|
||||
|
||||
try:
|
||||
# Check if we can use interactive mode
|
||||
if not sys.stdin.isatty():
|
||||
# Fallback to simple selection for non-interactive environments
|
||||
return prompt_team_selection_fallback(teams)
|
||||
|
||||
while True:
|
||||
display_interactive_team_selection(teams, selected_index)
|
||||
|
||||
key = get_key_input()
|
||||
|
||||
if key == "up":
|
||||
selected_index = (selected_index - 1) % len(teams)
|
||||
elif key == "down":
|
||||
selected_index = (selected_index + 1) % len(teams)
|
||||
elif key == "enter":
|
||||
selected_team = teams[selected_index]
|
||||
# Clear screen and show selection
|
||||
console = Console()
|
||||
console.clear()
|
||||
click.echo(
|
||||
f"✅ Selected team: {selected_team.get('team_alias', 'N/A')} ({selected_team.get('team_id')})"
|
||||
)
|
||||
return selected_team
|
||||
elif key == "quit" or key == "escape":
|
||||
# Clear screen
|
||||
console = Console()
|
||||
console.clear()
|
||||
click.echo("ℹ️ Team selection skipped.")
|
||||
return None
|
||||
elif key is None:
|
||||
# If we can't get key input, fall back to simple selection
|
||||
return prompt_team_selection_fallback(teams)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
console = Console()
|
||||
console.clear()
|
||||
click.echo("\n❌ Team selection cancelled.")
|
||||
return None
|
||||
except Exception:
|
||||
# If interactive mode fails, fall back to simple selection
|
||||
return prompt_team_selection_fallback(teams)
|
||||
|
||||
|
||||
def prompt_team_selection_fallback(
|
||||
teams: List[Dict[str, Any]]
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Fallback team selection for non-interactive environments"""
|
||||
if not teams:
|
||||
return None
|
||||
|
||||
while True:
|
||||
try:
|
||||
choice = click.prompt(
|
||||
"\nSelect a team by entering the index number (or 'skip' to continue without a team)",
|
||||
type=str,
|
||||
).strip()
|
||||
|
||||
if choice.lower() == "skip":
|
||||
return None
|
||||
|
||||
index = int(choice) - 1
|
||||
if 0 <= index < len(teams):
|
||||
selected_team = teams[index]
|
||||
click.echo(
|
||||
f"\n✅ Selected team: {selected_team.get('team_alias', 'N/A')} ({selected_team.get('team_id')})"
|
||||
)
|
||||
return selected_team
|
||||
else:
|
||||
click.echo(
|
||||
f"❌ Invalid selection. Please enter a number between 1 and {len(teams)}"
|
||||
)
|
||||
except ValueError:
|
||||
click.echo("❌ Invalid input. Please enter a number or 'skip'")
|
||||
except KeyboardInterrupt:
|
||||
click.echo("\n❌ Team selection cancelled.")
|
||||
return None
|
||||
|
||||
|
||||
# Polling-based authentication - no local server needed
|
||||
def _poll_for_ready_data(
|
||||
url: str,
|
||||
*,
|
||||
total_timeout: int = 300,
|
||||
poll_interval: int = 2,
|
||||
request_timeout: int = 10,
|
||||
pending_message: Optional[str] = None,
|
||||
pending_log_every: int = 10,
|
||||
other_status_message: Optional[str] = None,
|
||||
other_status_log_every: int = 10,
|
||||
http_error_log_every: int = 10,
|
||||
connection_error_log_every: int = 10,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
for attempt in range(total_timeout // poll_interval):
|
||||
try:
|
||||
response = requests.get(url, timeout=request_timeout)
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
status = data.get("status")
|
||||
if status == "ready":
|
||||
return data
|
||||
if status == "pending":
|
||||
if (
|
||||
pending_message
|
||||
and pending_log_every > 0
|
||||
and attempt % pending_log_every == 0
|
||||
):
|
||||
click.echo(pending_message)
|
||||
elif (
|
||||
other_status_message
|
||||
and other_status_log_every > 0
|
||||
and attempt % other_status_log_every == 0
|
||||
):
|
||||
click.echo(other_status_message)
|
||||
elif http_error_log_every > 0 and attempt % http_error_log_every == 0:
|
||||
click.echo(f"Polling error: HTTP {response.status_code}")
|
||||
except requests.RequestException as e:
|
||||
if (
|
||||
connection_error_log_every > 0
|
||||
and attempt % connection_error_log_every == 0
|
||||
):
|
||||
click.echo(f"Connection error (will retry): {e}")
|
||||
time.sleep(poll_interval)
|
||||
return None
|
||||
|
||||
|
||||
def _normalize_teams(teams, team_details):
|
||||
"""If team_details are a
|
||||
|
||||
Args:
|
||||
teams (_type_): _description_
|
||||
team_details (_type_): _description_
|
||||
|
||||
Returns:
|
||||
_type_: _description_
|
||||
"""
|
||||
if isinstance(team_details, list) and team_details:
|
||||
return [
|
||||
{
|
||||
"team_id": i.get("team_id") or i.get("id"),
|
||||
"team_alias": i.get("team_alias"),
|
||||
}
|
||||
for i in team_details
|
||||
if isinstance(i, dict) and (i.get("team_id") or i.get("id"))
|
||||
]
|
||||
if isinstance(teams, list):
|
||||
return [{"team_id": str(t), "team_alias": None} for t in teams]
|
||||
return []
|
||||
|
||||
|
||||
def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
|
||||
"""
|
||||
Poll the server for authentication completion and handle team selection.
|
||||
|
||||
Returns:
|
||||
Dictionary with authentication data if successful, None otherwise
|
||||
"""
|
||||
poll_url = f"{base_url}/sso/cli/poll/{key_id}"
|
||||
data = _poll_for_ready_data(
|
||||
poll_url,
|
||||
pending_message="Still waiting for authentication...",
|
||||
)
|
||||
if not data:
|
||||
return None
|
||||
if data.get("requires_team_selection"):
|
||||
teams = data.get("teams", [])
|
||||
team_details = data.get("team_details")
|
||||
user_id = data.get("user_id")
|
||||
normalized_teams: List[Dict[str, Any]] = _normalize_teams(teams, team_details)
|
||||
if not normalized_teams:
|
||||
click.echo("⚠️ No teams available for selection.")
|
||||
return None
|
||||
|
||||
# User has multiple teams - let them select
|
||||
jwt_with_team = _handle_team_selection_during_polling(
|
||||
base_url=base_url,
|
||||
key_id=key_id,
|
||||
teams=normalized_teams,
|
||||
)
|
||||
|
||||
# Use the team-specific JWT if selection succeeded
|
||||
if jwt_with_team:
|
||||
return {
|
||||
"api_key": jwt_with_team,
|
||||
"user_id": user_id,
|
||||
"teams": teams,
|
||||
"team_id": None, # Set by server in JWT
|
||||
}
|
||||
|
||||
click.echo("❌ Team selection cancelled or JWT generation failed.")
|
||||
return None
|
||||
|
||||
# JWT is ready (single team or team already selected)
|
||||
api_key = data.get("key")
|
||||
user_id = data.get("user_id")
|
||||
teams = data.get("teams", [])
|
||||
team_id = data.get("team_id")
|
||||
|
||||
# Show which team was assigned
|
||||
if team_id and len(teams) == 1:
|
||||
click.echo(f"\n✅ Automatically assigned to team: {team_id}")
|
||||
|
||||
if api_key:
|
||||
return {
|
||||
"api_key": api_key,
|
||||
"user_id": user_id,
|
||||
"teams": teams,
|
||||
"team_id": team_id,
|
||||
}
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _handle_team_selection_during_polling(
|
||||
base_url: str, key_id: str, teams: List[Dict[str, Any]]
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Handle team selection and re-poll with selected team_id.
|
||||
|
||||
Args:
|
||||
teams: List of team IDs (strings)
|
||||
|
||||
Returns:
|
||||
The JWT token with the selected team, or None if selection was skipped
|
||||
"""
|
||||
if not teams:
|
||||
click.echo(
|
||||
"ℹ️ No teams found. You can create or join teams using the web interface."
|
||||
)
|
||||
return None
|
||||
|
||||
click.echo("\n" + "=" * 60)
|
||||
click.echo("📋 Select a team for your CLI session...")
|
||||
|
||||
team_id = _render_and_prompt_for_team_selection(teams)
|
||||
|
||||
if not team_id:
|
||||
click.echo("ℹ️ No team selected.")
|
||||
return None
|
||||
|
||||
click.echo(f"\n🔄 Generating JWT for team: {team_id}")
|
||||
|
||||
poll_url = f"{base_url}/sso/cli/poll/{key_id}?team_id={team_id}"
|
||||
data = _poll_for_ready_data(
|
||||
poll_url,
|
||||
pending_message="Still waiting for team authentication...",
|
||||
other_status_message="Waiting for team authentication to complete...",
|
||||
http_error_log_every=10,
|
||||
)
|
||||
if not data:
|
||||
return None
|
||||
jwt_token = data.get("key")
|
||||
if jwt_token:
|
||||
click.echo(f"✅ Successfully generated JWT for team: {team_id}")
|
||||
return jwt_token
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _render_and_prompt_for_team_selection(teams: List[Dict[str, Any]]) -> Optional[str]:
|
||||
"""Render teams table and prompt user for a team selection.
|
||||
|
||||
Returns the selected team_id as a string, or None if selection was
|
||||
cancelled or skipped without any teams available.
|
||||
"""
|
||||
# Display teams as a simple list, but prefer showing aliases where
|
||||
# available while still keeping the underlying IDs intact.
|
||||
console = Console()
|
||||
table = Table(title="Available Teams")
|
||||
table.add_column("Index", style="cyan", no_wrap=True)
|
||||
table.add_column("Team Name", style="magenta")
|
||||
table.add_column("Team ID", style="green")
|
||||
|
||||
for i, team in enumerate(teams):
|
||||
team_id = str(team.get("team_id"))
|
||||
team_alias = team.get("team_alias") or team_id
|
||||
table.add_row(str(i + 1), team_alias, team_id)
|
||||
|
||||
console.print(table)
|
||||
|
||||
# Simple selection
|
||||
while True:
|
||||
try:
|
||||
choice = click.prompt(
|
||||
"\nSelect a team by entering the index number (or 'skip' to use first team)",
|
||||
type=str,
|
||||
).strip()
|
||||
|
||||
if choice.lower() == "skip":
|
||||
# Default to the first team's ID if the user skips an
|
||||
# explicit selection.
|
||||
if teams:
|
||||
first_team = teams[0]
|
||||
return str(first_team.get("team_id"))
|
||||
return None
|
||||
|
||||
index = int(choice) - 1
|
||||
if 0 <= index < len(teams):
|
||||
selected_team = teams[index]
|
||||
team_id = str(selected_team.get("team_id"))
|
||||
team_alias = selected_team.get("team_alias") or team_id
|
||||
click.echo(f"\n✅ Selected team: {team_alias} ({team_id})")
|
||||
return team_id
|
||||
|
||||
click.echo(
|
||||
f"❌ Invalid selection. Please enter a number between 1 and {len(teams)}"
|
||||
)
|
||||
except ValueError:
|
||||
click.echo("❌ Invalid input. Please enter a number or 'skip'")
|
||||
except KeyboardInterrupt:
|
||||
click.echo("\n❌ Team selection cancelled.")
|
||||
return None
|
||||
|
||||
|
||||
@click.command(name="login")
|
||||
@click.pass_context
|
||||
def login(ctx: click.Context):
|
||||
"""Login to LiteLLM proxy using SSO authentication"""
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER
|
||||
from litellm.proxy.client.cli.interface import show_commands
|
||||
|
||||
base_url = ctx.obj["base_url"]
|
||||
|
||||
# Check if we have an existing key to regenerate
|
||||
existing_key = get_stored_api_key()
|
||||
|
||||
# Generate unique key ID for this login session
|
||||
key_id = f"sk-{str(uuid.uuid4())}"
|
||||
|
||||
try:
|
||||
# Construct SSO login URL with CLI source and pre-generated key
|
||||
sso_url = f"{base_url}/sso/key/generate?source={LITELLM_CLI_SOURCE_IDENTIFIER}&key={key_id}"
|
||||
|
||||
# If we have an existing key, include it as a parameter to the login endpoint
|
||||
# The server will encode it in the OAuth state parameter for the SSO flow
|
||||
if existing_key:
|
||||
sso_url += f"&existing_key={existing_key}"
|
||||
|
||||
click.echo(f"Opening browser to: {sso_url}")
|
||||
click.echo("Please complete the SSO authentication in your browser...")
|
||||
click.echo(f"Session ID: {key_id}")
|
||||
|
||||
# Open browser
|
||||
webbrowser.open(sso_url)
|
||||
|
||||
# Poll for authentication completion
|
||||
click.echo("Waiting for authentication...")
|
||||
|
||||
auth_result = _poll_for_authentication(base_url=base_url, key_id=key_id)
|
||||
|
||||
if auth_result:
|
||||
api_key = auth_result["api_key"]
|
||||
user_id = auth_result["user_id"]
|
||||
|
||||
# Save token data (simplified for CLI - we just need the key)
|
||||
save_token(
|
||||
{
|
||||
"key": api_key,
|
||||
"user_id": user_id or "cli-user",
|
||||
"user_email": "unknown",
|
||||
"user_role": "cli",
|
||||
"auth_header_name": "Authorization",
|
||||
"jwt_token": "",
|
||||
"timestamp": time.time(),
|
||||
}
|
||||
)
|
||||
|
||||
click.echo("\n✅ Login successful!")
|
||||
click.echo(f"JWT Token: {api_key[:20]}...")
|
||||
click.echo("You can now use the CLI without specifying --api-key")
|
||||
|
||||
# Show available commands after successful login
|
||||
click.echo("\n" + "=" * 60)
|
||||
show_commands()
|
||||
return
|
||||
else:
|
||||
click.echo("❌ Authentication timed out. Please try again.")
|
||||
return
|
||||
|
||||
except KeyboardInterrupt:
|
||||
click.echo("\n❌ Authentication cancelled by user.")
|
||||
return
|
||||
except Exception as e:
|
||||
click.echo(f"❌ Authentication failed: {e}")
|
||||
return
|
||||
|
||||
|
||||
@click.command(name="logout")
|
||||
def logout():
|
||||
"""Logout and clear stored authentication"""
|
||||
clear_token()
|
||||
click.echo("✅ Logged out successfully. Authentication token cleared.")
|
||||
|
||||
|
||||
@click.command(name="whoami")
|
||||
def whoami():
|
||||
"""Show current authentication status"""
|
||||
token_data = load_token()
|
||||
|
||||
if not token_data:
|
||||
click.echo("❌ Not authenticated. Run 'litellm-proxy login' to authenticate.")
|
||||
return
|
||||
|
||||
click.echo("✅ Authenticated")
|
||||
click.echo(f"User Email: {token_data.get('user_email', 'Unknown')}")
|
||||
click.echo(f"User ID: {token_data.get('user_id', 'Unknown')}")
|
||||
click.echo(f"User Role: {token_data.get('user_role', 'Unknown')}")
|
||||
|
||||
# Check if token is still valid (basic timestamp check)
|
||||
timestamp = token_data.get("timestamp", 0)
|
||||
age_hours = (time.time() - timestamp) / 3600
|
||||
click.echo(f"Token age: {age_hours:.1f} hours")
|
||||
|
||||
if age_hours > CLI_JWT_EXPIRATION_HOURS:
|
||||
click.echo(
|
||||
f"⚠️ Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired."
|
||||
)
|
||||
|
||||
|
||||
# Export functions for use by other CLI commands
|
||||
__all__ = ["login", "logout", "whoami", "prompt_team_selection"]
|
||||
|
||||
# Export individual commands instead of grouping them
|
||||
# login, logout, and whoami will be added as top-level commands
|
||||
Reference in New Issue
Block a user