Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/client/cli/commands/keys.py: 0%
223 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1import builtins
2import json
3from collections.abc import Mapping, Sequence
4from datetime import datetime
5from typing import Any, Final, Literal
7import click
8import requests
9import rich
10from rich.table import Table
11from typing_extensions import ReadOnly, TypedDict
13from ...keys import KeysManagementClient
16class _CliContext(TypedDict):
17 """Values the top-level CLI group stores on the click context."""
19 base_url: ReadOnly[str]
20 api_key: ReadOnly[str | None]
23class _CliContextView(TypedDict):
24 obj: ReadOnly[_CliContext]
27class _KeyRowsView(TypedDict):
28 rows: ReadOnly[Sequence[Mapping[str, object]]]
31class _JsonBodyView(TypedDict):
32 body: ReadOnly[object]
35@click.group()
36def keys():
37 """Manage API keys for the LiteLLM proxy server"""
40@keys.command()
41@click.option("--page", type=int, help="Page number for pagination")
42@click.option("--size", type=int, help="Number of items per page")
43@click.option("--user-id", type=str, help="Filter keys by user ID")
44@click.option("--team-id", type=str, help="Filter keys by team ID")
45@click.option("--organization-id", type=str, help="Filter keys by organization ID")
46@click.option("--key-hash", type=str, help="Filter by specific key hash")
47@click.option("--key-alias", type=str, help="Filter by key alias")
48@click.option(
49 "--return-full-object",
50 is_flag=True,
51 default=True,
52 help="Return the full key object",
53)
54@click.option("--include-team-keys", is_flag=True, help="Include team keys in the response")
55@click.option(
56 "--format",
57 "output_format",
58 type=click.Choice(["table", "json"]),
59 default="table",
60 help="Output format (table or json)",
61)
62@click.pass_context
63def list(
64 ctx: click.Context,
65 page: int | None,
66 size: int | None,
67 user_id: str | None,
68 team_id: str | None,
69 organization_id: str | None,
70 key_hash: str | None,
71 key_alias: str | None,
72 include_team_keys: bool,
73 output_format: Literal["table", "json"],
74 return_full_object: bool,
75):
76 """List all API keys"""
77 context: Final[_CliContextView] = {"obj": ctx.obj}
78 client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
79 response: Final = client.list(
80 page=page,
81 size=size,
82 user_id=user_id,
83 team_id=team_id,
84 organization_id=organization_id,
85 key_hash=key_hash,
86 key_alias=key_alias,
87 return_full_object=return_full_object,
88 include_team_keys=include_team_keys,
89 )
90 assert isinstance(response, dict)
92 if output_format == "json":
93 rich.print_json(data=response)
94 else:
95 listed: Final[_KeyRowsView] = {"rows": response.get("keys", [])}
96 rich.print(f"Showing {len(listed['rows'])} keys out of {response.get('total_count', 0)}")
97 table: Final = Table(title="API Keys")
98 table.add_column("Key Hash", style="cyan")
99 table.add_column("Alias", style="green")
100 table.add_column("User ID", style="magenta")
101 table.add_column("Team ID", style="yellow")
102 table.add_column("Spend", style="red")
103 key_rows: Final[_KeyRowsView] = {"rows": response.get("keys", [])}
104 for key in key_rows["rows"]:
105 table.add_row(
106 str(key.get("token", "")),
107 str(key.get("key_alias", "")),
108 str(key.get("user_id", "")),
109 str(key.get("team_id", "")),
110 str(key.get("spend", "")),
111 )
112 rich.print(table)
115@keys.command()
116@click.option("--models", type=str, help="Comma-separated list of allowed models")
117@click.option("--aliases", type=str, help="JSON string of model alias mappings")
118@click.option("--spend", type=float, help="Maximum spend limit for this key")
119@click.option(
120 "--duration",
121 type=str,
122 help="Duration for which the key is valid (e.g. '24h', '7d')",
123)
124@click.option("--key-alias", type=str, help="Alias/name for the key")
125@click.option("--team-id", type=str, help="Team ID to associate the key with")
126@click.option("--user-id", type=str, help="User ID to associate the key with")
127@click.option("--budget-id", type=str, help="Budget ID to associate the key with")
128@click.option("--config", type=str, help="JSON string of additional configuration parameters")
129@click.pass_context
130def generate(
131 ctx: click.Context,
132 models: str | None,
133 aliases: str | None,
134 spend: float | None,
135 duration: str | None,
136 key_alias: str | None,
137 team_id: str | None,
138 user_id: str | None,
139 budget_id: str | None,
140 config: str | None,
141):
142 """Generate a new API key"""
143 context: Final[_CliContextView] = {"obj": ctx.obj}
144 client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
145 try:
146 models_list: Final = [m.strip() for m in models.split(",")] if models else None
147 aliases_dict: Final = json.loads(aliases) if aliases else None
148 config_dict: Final = json.loads(config) if config else None
149 except json.JSONDecodeError as e:
150 raise click.BadParameter(f"Invalid JSON: {e}")
151 try:
152 response: Final = client.generate(
153 models=models_list,
154 aliases=aliases_dict,
155 spend=spend,
156 duration=duration,
157 key_alias=key_alias,
158 team_id=team_id,
159 user_id=user_id,
160 budget_id=budget_id,
161 config=config_dict,
162 )
163 rich.print_json(data=response)
164 except requests.exceptions.HTTPError as e:
165 click.echo(f"Error: HTTP {e.response.status_code}", err=True)
166 try:
167 error_body: Final[_JsonBodyView] = {"body": e.response.json()}
168 rich.print_json(data=error_body["body"])
169 except json.JSONDecodeError:
170 click.echo(e.response.text, err=True)
171 raise click.Abort()
174@keys.command()
175@click.option("--keys", type=str, help="Comma-separated list of API keys to delete")
176@click.option("--key-aliases", type=str, help="Comma-separated list of key aliases to delete")
177@click.pass_context
178def delete(ctx: click.Context, keys: str | None, key_aliases: str | None):
179 """Delete API keys by key or alias"""
180 context: Final[_CliContextView] = {"obj": ctx.obj}
181 client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
182 keys_list: Final = [k.strip() for k in keys.split(",")] if keys else None
183 aliases_list: Final = [a.strip() for a in key_aliases.split(",")] if key_aliases else None
184 try:
185 response: Final = client.delete(keys=keys_list, key_aliases=aliases_list)
186 rich.print_json(data=response)
187 except requests.exceptions.HTTPError as e:
188 click.echo(f"Error: HTTP {e.response.status_code}", err=True)
189 try:
190 error_body: Final[_JsonBodyView] = {"body": e.response.json()}
191 rich.print_json(data=error_body["body"])
192 except json.JSONDecodeError:
193 click.echo(e.response.text, err=True)
194 raise click.Abort()
197def _parse_created_since_filter(created_since: str | None) -> datetime | None:
198 """Parse and validate the created_since date filter."""
199 if not created_since:
200 return None
202 try:
203 # Support formats: YYYY-MM-DD_HH:MM or YYYY-MM-DD
204 if "_" in created_since:
205 return datetime.strptime(created_since, "%Y-%m-%d_%H:%M")
206 else:
207 return datetime.strptime(created_since, "%Y-%m-%d")
208 except ValueError:
209 click.echo(
210 f"Error: Invalid date format '{created_since}'. Use YYYY-MM-DD_HH:MM or YYYY-MM-DD",
211 err=True,
212 )
213 raise click.Abort()
216def _fetch_all_keys_with_pagination(
217 source_client: KeysManagementClient, source_base_url: str
218) -> Sequence[Mapping[str, object]]:
219 """Fetch all keys from source instance using pagination."""
220 click.echo(f"Fetching keys from source server: {source_base_url}")
221 source_keys: Final[builtins.list[Mapping[str, object]]] = []
222 page = 1
223 page_size: Final = 100 # Use a larger page size to minimize API calls
225 while True:
226 source_response = source_client.list(return_full_object=True, page=page, size=page_size)
227 # source_client.list() returns Dict[str, Any] when return_request is False (default)
228 assert isinstance(source_response, dict), "Expected dict response from list API"
229 page_keys: Sequence[Mapping[str, object]] = source_response.get("keys", [])
231 if not page_keys:
232 break
234 source_keys.extend(page_keys)
235 click.echo(f"Fetched page {page}: {len(page_keys)} keys")
237 # Check if we got fewer keys than the page size, indicating last page
238 if len(page_keys) < page_size:
239 break
241 page += 1
243 return source_keys
246def _filter_keys_by_created_since(
247 source_keys: Sequence[Mapping[str, object]],
248 created_since_dt: datetime | None,
249 created_since: str,
250) -> Sequence[Mapping[str, object]]:
251 """Filter keys by created_since date if specified."""
252 if not created_since_dt:
253 return source_keys
255 filtered_keys: Final[builtins.list[Mapping[str, object]]] = []
256 for key in source_keys:
257 key_created_at = key.get("created_at")
258 if key_created_at:
259 # Parse the key's created_at timestamp
260 if isinstance(key_created_at, str):
261 if "T" in key_created_at:
262 key_dt = datetime.fromisoformat(key_created_at.replace("Z", "+00:00"))
263 else:
264 key_dt = datetime.fromisoformat(key_created_at)
266 # Convert to naive datetime for comparison (assuming UTC)
267 if key_dt.tzinfo:
268 key_dt = key_dt.replace(tzinfo=None)
270 if key_dt >= created_since_dt:
271 filtered_keys.append(key)
273 click.echo(f"Filtered {len(source_keys)} keys to {len(filtered_keys)} keys created since {created_since}")
274 return filtered_keys
277def _display_dry_run_table(source_keys: Sequence[Mapping[str, object]]) -> None:
278 """Display a table of keys that would be imported in dry-run mode."""
279 click.echo("\n--- DRY RUN MODE ---")
280 table: Final = Table(title="Keys that would be imported")
281 table.add_column("Key Alias", style="green")
282 table.add_column("User ID", style="magenta")
283 table.add_column("Created", style="cyan")
285 for key in source_keys:
286 created_at = key.get("created_at", "")
287 # Format the timestamp if it exists
288 if created_at:
289 # Try to parse and format the timestamp for better readability
290 if isinstance(created_at, str):
291 # Handle common timestamp formats
292 if "T" in created_at:
293 dt = datetime.fromisoformat(created_at.replace("Z", "+00:00"))
294 created_at = dt.strftime("%Y-%m-%d %H:%M")
296 table.add_row(str(key.get("key_alias", "")), str(key.get("user_id", "")), str(created_at))
297 rich.print(table)
300def _prepare_key_import_data(key: Mapping[str, object]) -> dict[str, Any]:
301 """Prepare key data for import by extracting relevant fields."""
302 import_data: Final = {}
304 # Copy relevant fields if they exist
305 for field in [
306 "models",
307 "aliases",
308 "spend",
309 "key_alias",
310 "team_id",
311 "user_id",
312 "budget_id",
313 "config",
314 ]:
315 if key.get(field):
316 import_data[field] = key[field]
318 return import_data
321def _import_keys_to_destination(
322 source_keys: Sequence[Mapping[str, object]], dest_client: KeysManagementClient
323) -> tuple[int, int]:
324 """Import each key to the destination instance and return counts."""
325 imported_count = 0
326 failed_count = 0
328 for key in source_keys:
329 try:
330 # Prepare key data for import
331 import_data = _prepare_key_import_data(key)
333 # Generate the key in destination instance
334 response = dest_client.generate(**import_data)
335 click.echo(f"Generated key: {response}")
336 # The generate method returns JSON data directly, not a Response object
337 imported_count += 1
339 key_alias = key.get("key_alias", "N/A")
340 click.echo(f"Imported key: {key_alias}")
342 except Exception as e:
343 failed_count += 1
344 key_alias = key.get("key_alias", "N/A")
345 click.echo(f"Failed to import key {key_alias}: {e}", err=True)
347 return imported_count, failed_count
350@keys.command(name="import")
351@click.option(
352 "--source-base-url",
353 required=True,
354 help="Base URL of the source LiteLLM proxy server to import keys from",
355)
356@click.option("--source-api-key", help="API key for authentication to the source server")
357@click.option(
358 "--dry-run",
359 is_flag=True,
360 help="Show what would be imported without actually importing",
361)
362@click.option(
363 "--created-since",
364 help="Only import keys created after this date/time (format: YYYY-MM-DD_HH:MM or YYYY-MM-DD)",
365)
366@click.pass_context
367def import_keys(
368 ctx: click.Context,
369 source_base_url: str,
370 source_api_key: str | None,
371 dry_run: bool,
372 created_since: str | None,
373):
374 """Import API keys from another LiteLLM instance"""
375 # Parse created_since filter if provided
376 created_since_dt: Final = _parse_created_since_filter(created_since)
378 # Create clients for both source and destination
379 source_client: Final = KeysManagementClient(source_base_url, source_api_key)
380 context: Final[_CliContextView] = {"obj": ctx.obj}
381 dest_client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
383 try:
384 # Get all keys from source instance with pagination
385 source_keys = _fetch_all_keys_with_pagination(source_client, source_base_url)
387 # Filter keys by created_since if specified
388 if created_since:
389 source_keys = _filter_keys_by_created_since(source_keys, created_since_dt, created_since)
391 if not source_keys:
392 click.echo("No keys found in source instance.")
393 return
395 click.echo(f"Found {len(source_keys)} keys in source instance.")
397 if dry_run:
398 _display_dry_run_table(source_keys)
399 return
401 # Import each key
402 imported_count, failed_count = _import_keys_to_destination(source_keys, dest_client)
404 # Summary
405 click.echo("\nImport completed:")
406 click.echo(f" Successfully imported: {imported_count}")
407 click.echo(f" Failed to import: {failed_count}")
408 click.echo(f" Total keys processed: {len(source_keys)}")
410 except requests.exceptions.HTTPError as e:
411 click.echo(f"Error: HTTP {e.response.status_code}", err=True)
412 try:
413 error_body: Final[_JsonBodyView] = {"body": e.response.json()}
414 rich.print_json(data=error_body["body"])
415 except json.JSONDecodeError:
416 click.echo(e.response.text, err=True)
417 raise click.Abort()
418 except Exception as e:
419 click.echo(f"Error: {e}", err=True)
420 raise click.Abort()