Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/client/cli/commands/teams.py: 0%

1"""Team management commands for LiteLLM CLI.""" 

2 

3from collections.abc import Mapping, Sequence 

4from typing import Any, Final 

5 

6import click 

7import requests 

8from rich.console import Console 

9from rich.table import Table 

10from typing_extensions import ReadOnly, TypedDict 

11 

12from litellm.proxy.client import Client 

13 

14from ._cli_context import cli_context_values 

15 

16 

17class _TeamRow(TypedDict): 

18 team_alias: ReadOnly[str | None] 

19 team_id: ReadOnly[str | None] 

20 models: ReadOnly[Sequence[str]] 

21 max_budget: ReadOnly[object] 

22 

23 

24class _TeamModelsView(TypedDict): 

25 models: ReadOnly[Sequence[str]] 

26 

27 

28class _ErrorBodyView(TypedDict): 

29 body: ReadOnly[Mapping[str, object]] 

30 

31 

32@click.group() 

33def teams(): 

34 """Manage teams and team assignments""" 

35 

36 

37def display_teams_table(teams: Sequence[dict[str, Any]]) -> None: 

38 """Display teams in a formatted table""" 

39 console: Final = Console() 

40 

41 if not teams: 

42 console.print("No teams found for your user.") 

43 return 

44 

45 table: Final = Table(title="Available Teams") 

46 table.add_column("Index", style="cyan", no_wrap=True) 

47 table.add_column("Team Alias", style="magenta") 

48 table.add_column("Team ID", style="green") 

49 table.add_column("Models", style="yellow") 

50 table.add_column("Max Budget", style="blue") 

51 table.add_column("Role", style="red") 

52 

53 for i, team in enumerate(teams): 

54 row: _TeamRow = { 

55 "team_alias": team.get("team_alias") or "N/A", 

56 "team_id": team.get("team_id", "N/A"), 

57 "models": team.get("models", []), 

58 "max_budget": team.get("max_budget"), 

59 } 

60 models = row["models"] 

61 max_budget = row["max_budget"] 

62 

63 # Format models list 

64 if models: 

65 if len(models) > 3: 

66 models_str = ", ".join(models[:3]) + f" (+{len(models) - 3} more)" 

67 else: 

68 models_str = ", ".join(models) 

69 else: 

70 models_str = "All models" 

71 

72 # Format budget 

73 budget_str = f"${max_budget}" if max_budget else "Unlimited" 

74 

75 # Try to determine role (this might vary based on API response structure) 

76 role = "Member" # Default role 

77 if isinstance(team, dict) and "members_with_roles" in team and team["members_with_roles"]: 

78 # This would need to be implemented based on actual API response structure 

79 pass 

80 

81 table.add_row(str(i + 1), row["team_alias"], row["team_id"], models_str, budget_str, role) 

82 

83 console.print(table) 

84 

85 

86@teams.command() 

87@click.pass_context 

88def list(ctx: click.Context): 

89 """List teams that you belong to""" 

90 context: Final = cli_context_values(ctx) 

91 client: Final = Client(context["base_url"], context["api_key"]) 

92 

93 try: 

94 # Use list() for simpler response structure (returns array directly) 

95 teams: Final = client.teams.list() 

96 display_teams_table(teams) 

97 except requests.exceptions.HTTPError as e: 

98 click.echo(f"Error: HTTP {e.response.status_code}", err=True) 

99 error_body: Final[_ErrorBodyView] = {"body": e.response.json()} 

100 click.echo(f"Details: {error_body['body'].get('detail', 'Unknown error')}", err=True) 

101 raise click.Abort() 

102 except Exception as e: 

103 click.echo(f"Error: {e}", err=True) 

104 raise click.Abort() 

105 

106 

107@teams.command() 

108@click.pass_context 

109def available(ctx: click.Context): 

110 """List teams that are available to join""" 

111 context: Final = cli_context_values(ctx) 

112 client: Final = Client(context["base_url"], context["api_key"]) 

113 

114 try: 

115 teams: Final = client.teams.get_available() 

116 if teams: 

117 console: Final = Console() 

118 console.print("\nAvailable Teams to Join:") 

119 display_teams_table(teams) 

120 else: 

121 click.echo("No available teams to join.") 

122 except requests.exceptions.HTTPError as e: 

123 click.echo(f"Error: HTTP {e.response.status_code}", err=True) 

124 error_body: Final[_ErrorBodyView] = {"body": e.response.json()} 

125 click.echo(f"Details: {error_body['body'].get('detail', 'Unknown error')}", err=True) 

126 except Exception as e: 

127 click.echo(f"Error: {e}", err=True) 

128 raise click.Abort() 

129 

130 

131@teams.command() 

132@click.option("--team-id", type=str, help="Team ID to assign the key to") 

133@click.pass_context 

134def assign_key(ctx: click.Context, team_id: str | None): 

135 """Assign your current CLI key to a team""" 

136 context: Final = cli_context_values(ctx) 

137 client: Final = Client(context["base_url"], context["api_key"]) 

138 api_key: Final = context["api_key"] 

139 

140 if not api_key: 

141 click.echo("No API key found. Please login first using 'litellm login'") 

142 raise click.Abort() 

143 

144 try: 

145 # If no team_id provided, show teams and let user select 

146 if not team_id: 

147 teams = client.teams.list() 

148 

149 if not teams: 

150 click.echo("No teams found for your user.") 

151 return 

152 

153 # Use interactive selection from auth module 

154 from .auth import prompt_team_selection 

155 

156 selected_team: Final = prompt_team_selection(teams) 

157 

158 if selected_team: 

159 team_id = selected_team.get("team_id") 

160 else: 

161 click.echo("Operation cancelled.") 

162 return 

163 

164 # Update the key with the selected team 

165 if team_id: 

166 click.echo(f"\nAssigning your key to team: {team_id}") 

167 client.keys.update(key=api_key, team_id=team_id) 

168 click.echo(f"Successfully assigned key to team: {team_id}") 

169 

170 # Show team details if available 

171 teams = client.teams.list() 

172 for team in teams: 

173 if team.get("team_id") == team_id: 

174 team_models: _TeamModelsView = {"models": team.get("models", [])} 

175 if team_models["models"]: 

176 click.echo(f"You can now access models: {', '.join(team_models['models'])}") 

177 else: 

178 click.echo("You can now access all available models") 

179 break 

180 

181 except requests.exceptions.HTTPError as e: 

182 click.echo(f"Error: HTTP {e.response.status_code}", err=True) 

183 error_body: Final[_ErrorBodyView] = {"body": e.response.json()} 

184 click.echo(f"Details: {error_body['body'].get('detail', 'Unknown error')}", err=True) 

185 raise click.Abort() 

186 except Exception as e: 

187 click.echo(f"Error: {e}", err=True) 

188 raise click.Abort()