diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 9be7426805..ae795316a2 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -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) diff --git a/litellm/proxy/client/cli/commands/teams.py b/litellm/proxy/client/cli/commands/teams.py index d4e05c890d..57397ca01a 100644 --- a/litellm/proxy/client/cli/commands/teams.py +++ b/litellm/proxy/client/cli/commands/teams.py @@ -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)