tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-05 08:23:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor

Browse Source
This commit is contained in:
Ishaan Jaffer
2025-09-23 14:07:33 -07:00
parent 298b0a50a7
commit d91e533d8e
2 changed files with 40 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/client/cli/commands/auth.py
+33 -29
View File
@@ -301,6 +301,37 @@ def update_key_with_team(base_url: str, api_key: str, team_id: str) -> bool:
# Polling-based authentication - no local server needed
def _handle_team_assignment(base_url: str, api_key: str, user_id: str) -> None:
"""Handle team fetching and assignment for the authenticated user."""
click.echo("\n" + "="*60)
click.echo("📋 Fetching your teams...")
teams = get_user_teams(
base_url=base_url,
api_key=api_key,
user_id=user_id,
)
if teams:
# Prompt for team selection (will display teams interactively)
selected_team = prompt_team_selection(teams)
if selected_team:
team_id = selected_team.get('team_id')
if team_id:
click.echo(f"\n🔄 Assigning your key to team: {selected_team.get('team_alias', team_id)}")
success = update_key_with_team(base_url, api_key, team_id)
if success:
click.echo(f"✅ Your CLI key is now associated with team: {selected_team.get('team_alias', team_id)}")
click.echo(f"🎯 You can now access models: {', '.join(selected_team.get('models', ['All models']))}")
else:
click.echo("⚠️ Key assignment failed, but you can still use the CLI")
else:
click.echo("ℹ️ Continuing without team assignment. You can assign a team later using the CLI.")
else:
click.echo("ℹ️ No teams found. You can create or join teams using the web interface.")
@click.command(name="login")
@click.pass_context
def login(ctx: click.Context):
@@ -364,35 +395,8 @@ def login(ctx: click.Context):
click.echo(f"API Key: {api_key[:20]}...")
click.echo("You can now use the CLI without specifying --api-key")
# Fetch and display user's teams
click.echo("\n" + "="*60)
click.echo("📋 Fetching your teams...")
teams = get_user_teams(
base_url=base_url,
api_key=api_key,
user_id=data.get("user_id"),
)
if teams:
# Prompt for team selection (will display teams interactively)
selected_team = prompt_team_selection(teams)
if selected_team:
team_id = selected_team.get('team_id')
if team_id:
click.echo(f"\n🔄 Assigning your key to team: {selected_team.get('team_alias', team_id)}")
success = update_key_with_team(base_url, api_key, team_id)
if success:
click.echo(f"✅ Your CLI key is now associated with team: {selected_team.get('team_alias', team_id)}")
click.echo(f"🎯 You can now access models: {', '.join(selected_team.get('models', ['All models']))}")
else:
click.echo("⚠️ Key assignment failed, but you can still use the CLI")
else:
click.echo("ℹ️ Continuing without team assignment. You can assign a team later using the CLI.")
else:
click.echo("ℹ️ No teams found. You can create or join teams using the web interface.")
# Handle team assignment
_handle_team_assignment(base_url, api_key, data.get("user_id"))
# Show available commands after successful login
click.echo("\n" + "="*60)
litellm/proxy/client/cli/commands/teams.py
+7 -17
View File
@@ -80,11 +80,8 @@ def list(ctx: click.Context):
display_teams_table(teams)
except requests.exceptions.HTTPError as e:
click.echo(f"Error: HTTP {e.response.status_code}", err=True)
try:
error_body = e.response.json()
click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True)
except:
click.echo(e.response.text, err=True)
error_body = e.response.json()
click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True)
raise click.Abort()
except Exception as e:
click.echo(f"Error: {str(e)}", err=True)
@@ -107,12 +104,8 @@ def available(ctx: click.Context):
click.echo("ℹ️ No available teams to join.")
except requests.exceptions.HTTPError as e:
click.echo(f"Error: HTTP {e.response.status_code}", err=True)
try:
error_body = e.response.json()
click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True)
except:
click.echo(e.response.text, err=True)
raise click.Abort()
error_body = e.response.json()
click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True)
except Exception as e:
click.echo(f"Error: {str(e)}", err=True)
raise click.Abort()
@@ -152,7 +145,7 @@ def assign_key(ctx: click.Context, team_id: Optional[str]):
# Update the key with the selected team
if team_id:
click.echo(f"\n🔄 Assigning your key to team: {team_id}")
result = client.keys.update(key=api_key, team_id=team_id)
client.keys.update(key=api_key, team_id=team_id)
click.echo(f"✅ Successfully assigned key to team: {team_id}")
# Show team details if available
@@ -168,11 +161,8 @@ def assign_key(ctx: click.Context, team_id: Optional[str]):
except requests.exceptions.HTTPError as e:
click.echo(f"Error: HTTP {e.response.status_code}", err=True)
try:
error_body = e.response.json()
click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True)
except:
click.echo(e.response.text, err=True)
error_body = e.response.json()
click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True)
raise click.Abort()
except Exception as e:
click.echo(f"Error: {str(e)}", err=True)