From 40ccd2b70f868ee403844b6ab8dff0caad35060e Mon Sep 17 00:00:00 2001 From: Marc Abramowitz Date: Tue, 15 Jul 2025 20:14:43 -0700 Subject: [PATCH] Add "keys import" command to CLI (#12620) * Add "keys import" command to CLI E.g.: ``` litellm-proxy keys import \ --source-base-url=https://old-litellm.company.com \ --source-api-key=$LITELLM_KEY \ --dry-run ``` * Add --created-since option * Add tests * Fix lint errors * Fix lint issues * Fix lint errors * Fix response.raise_for_status not being a thing * Fix a mypy error --- litellm/proxy/client/cli/commands/keys.py | 201 ++++++++++- .../proxy/client/cli/test_keys_commands.py | 333 +++++++++++++++++- 2 files changed, 527 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/client/cli/commands/keys.py b/litellm/proxy/client/cli/commands/keys.py index 5ae81fb24d..2e10304b3f 100644 --- a/litellm/proxy/client/cli/commands/keys.py +++ b/litellm/proxy/client/cli/commands/keys.py @@ -1,5 +1,6 @@ import json -from typing import Literal, Optional +from datetime import datetime +from typing import Literal, Optional, List, Dict, Any import click import rich @@ -156,3 +157,201 @@ def delete(ctx: click.Context, keys: Optional[str], key_aliases: Optional[str]): except json.JSONDecodeError: click.echo(e.response.text, err=True) raise click.Abort() + + +def _parse_created_since_filter(created_since: Optional[str]) -> Optional[datetime]: + """Parse and validate the created_since date filter.""" + if not created_since: + return None + + try: + # Support formats: YYYY-MM-DD_HH:MM or YYYY-MM-DD + if "_" in created_since: + return datetime.strptime(created_since, "%Y-%m-%d_%H:%M") + else: + return datetime.strptime(created_since, "%Y-%m-%d") + except ValueError: + click.echo(f"Error: Invalid date format '{created_since}'. Use YYYY-MM-DD_HH:MM or YYYY-MM-DD", err=True) + raise click.Abort() + + +def _fetch_all_keys_with_pagination(source_client: KeysManagementClient, source_base_url: str) -> List[Dict[str, Any]]: + """Fetch all keys from source instance using pagination.""" + click.echo(f"Fetching keys from source server: {source_base_url}") + source_keys = [] + page = 1 + page_size = 100 # Use a larger page size to minimize API calls + + while True: + source_response = source_client.list(return_full_object=True, page=page, size=page_size) + # source_client.list() returns Dict[str, Any] when return_request is False (default) + assert isinstance(source_response, dict), "Expected dict response from list API" + page_keys = source_response.get("keys", []) + + if not page_keys: + break + + source_keys.extend(page_keys) + click.echo(f"Fetched page {page}: {len(page_keys)} keys") + + # Check if we got fewer keys than the page size, indicating last page + if len(page_keys) < page_size: + break + + page += 1 + + return source_keys + + +def _filter_keys_by_created_since( + source_keys: List[Dict[str, Any]], created_since_dt: Optional[datetime], created_since: str +) -> List[Dict[str, Any]]: + """Filter keys by created_since date if specified.""" + if not created_since_dt: + return source_keys + + filtered_keys = [] + for key in source_keys: + key_created_at = key.get("created_at") + if key_created_at: + # Parse the key's created_at timestamp + if isinstance(key_created_at, str): + if "T" in key_created_at: + key_dt = datetime.fromisoformat(key_created_at.replace("Z", "+00:00")) + else: + key_dt = datetime.fromisoformat(key_created_at) + + # Convert to naive datetime for comparison (assuming UTC) + if key_dt.tzinfo: + key_dt = key_dt.replace(tzinfo=None) + + if key_dt >= created_since_dt: + filtered_keys.append(key) + + click.echo(f"Filtered {len(source_keys)} keys to {len(filtered_keys)} keys created since {created_since}") + return filtered_keys + + +def _display_dry_run_table(source_keys: List[Dict[str, Any]]) -> None: + """Display a table of keys that would be imported in dry-run mode.""" + click.echo("\n--- DRY RUN MODE ---") + table = Table(title="Keys that would be imported") + table.add_column("Key Alias", style="green") + table.add_column("User ID", style="magenta") + table.add_column("Created", style="cyan") + + for key in source_keys: + created_at = key.get("created_at", "") + # Format the timestamp if it exists + if created_at: + # Try to parse and format the timestamp for better readability + if isinstance(created_at, str): + # Handle common timestamp formats + if "T" in created_at: + dt = datetime.fromisoformat(created_at.replace("Z", "+00:00")) + created_at = dt.strftime("%Y-%m-%d %H:%M") + + table.add_row(str(key.get("key_alias", "")), str(key.get("user_id", "")), str(created_at)) + rich.print(table) + + +def _prepare_key_import_data(key: Dict[str, Any]) -> Dict[str, Any]: + """Prepare key data for import by extracting relevant fields.""" + import_data = {} + + # Copy relevant fields if they exist + for field in ["models", "aliases", "spend", "key_alias", "team_id", "user_id", "budget_id", "config"]: + if key.get(field): + import_data[field] = key[field] + + return import_data + + +def _import_keys_to_destination( + source_keys: List[Dict[str, Any]], dest_client: KeysManagementClient +) -> tuple[int, int]: + """Import each key to the destination instance and return counts.""" + imported_count = 0 + failed_count = 0 + + for key in source_keys: + try: + # Prepare key data for import + import_data = _prepare_key_import_data(key) + + # Generate the key in destination instance + response = dest_client.generate(**import_data) + click.echo(f"Generated key: {response}") + # The generate method returns JSON data directly, not a Response object + imported_count += 1 + + key_alias = key.get("key_alias", "N/A") + click.echo(f"✓ Imported key: {key_alias}") + + except Exception as e: + failed_count += 1 + key_alias = key.get("key_alias", "N/A") + click.echo(f"✗ Failed to import key {key_alias}: {str(e)}", err=True) + + return imported_count, failed_count + + +@keys.command(name="import") +@click.option( + "--source-base-url", required=True, help="Base URL of the source LiteLLM proxy server to import keys from" +) +@click.option("--source-api-key", help="API key for authentication to the source server") +@click.option("--dry-run", is_flag=True, help="Show what would be imported without actually importing") +@click.option( + "--created-since", help="Only import keys created after this date/time (format: YYYY-MM-DD_HH:MM or YYYY-MM-DD)" +) +@click.pass_context +def import_keys( + ctx: click.Context, source_base_url: str, source_api_key: Optional[str], dry_run: bool, created_since: Optional[str] +): + """Import API keys from another LiteLLM instance""" + # Parse created_since filter if provided + created_since_dt = _parse_created_since_filter(created_since) + + # Create clients for both source and destination + source_client = KeysManagementClient(source_base_url, source_api_key) + dest_client = KeysManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + + try: + # Get all keys from source instance with pagination + source_keys = _fetch_all_keys_with_pagination(source_client, source_base_url) + + # Filter keys by created_since if specified + if created_since: + source_keys = _filter_keys_by_created_since(source_keys, created_since_dt, created_since) + + if not source_keys: + click.echo("No keys found in source instance.") + return + + click.echo(f"Found {len(source_keys)} keys in source instance.") + + if dry_run: + _display_dry_run_table(source_keys) + return + + # Import each key + imported_count, failed_count = _import_keys_to_destination(source_keys, dest_client) + + # Summary + click.echo("\nImport completed:") + click.echo(f" Successfully imported: {imported_count}") + click.echo(f" Failed to import: {failed_count}") + click.echo(f" Total keys processed: {len(source_keys)}") + + except requests.exceptions.HTTPError as e: + click.echo(f"Error: HTTP {e.response.status_code}", err=True) + try: + error_body = e.response.json() + rich.print_json(data=error_body) + except json.JSONDecodeError: + click.echo(e.response.text, err=True) + raise click.Abort() + except Exception as e: + click.echo(f"Error: {str(e)}", err=True) + raise click.Abort() diff --git a/tests/test_litellm/proxy/client/cli/test_keys_commands.py b/tests/test_litellm/proxy/client/cli/test_keys_commands.py index 8b7abc4ea4..423e23400f 100644 --- a/tests/test_litellm/proxy/client/cli/test_keys_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_keys_commands.py @@ -3,6 +3,8 @@ import os import sys from unittest.mock import patch +import requests + sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path @@ -34,9 +36,7 @@ def mock_env(): @pytest.fixture def mock_keys_client(): - with patch( - "litellm.proxy.client.cli.commands.keys.KeysManagementClient" - ) as MockClient: + with patch("litellm.proxy.client.cli.commands.keys.KeysManagementClient") as MockClient: yield MockClient @@ -88,9 +88,7 @@ def test_async_keys_generate_success(mock_keys_client, cli_runner): "key": "new-key", "spend": 100.0, } - result = cli_runner.invoke( - cli, ["keys", "generate", "--models", "gpt-4", "--spend", "100"] - ) + result = cli_runner.invoke(cli, ["keys", "generate", "--models", "gpt-4", "--spend", "100"]) assert result.exit_code == 0 assert "new-key" in result.output mock_keys_client.return_value.generate.assert_called_once() @@ -158,3 +156,326 @@ def test_async_keys_delete_http_error_handling(mock_keys_client, cli_runner): assert result.exit_code != 0 # HTTPError should be caught and converted to click.Abort assert isinstance(result.exception, SystemExit) # click.Abort raises SystemExit + + +# Tests for keys import command +def test_keys_import_dry_run_success(mock_keys_client, cli_runner): + """Test successful dry-run import showing table of keys that would be imported""" + # Mock source client response (paginated) + mock_source_instance = mock_keys_client.return_value + mock_source_instance.list.side_effect = [ + { + "keys": [ + { + "key_alias": "test-key-1", + "user_id": "user1@example.com", + "created_at": "2024-01-15T10:30:00Z", + "models": ["gpt-4"], + "spend": 10.0, + }, + { + "key_alias": "test-key-2", + "user_id": "user2@example.com", + "created_at": "2024-01-16T11:45:00Z", + "models": [], + "spend": 5.0, + } + ] + }, + {"keys": []} # Empty second page + ] + + result = cli_runner.invoke(cli, [ + "keys", "import", + "--source-base-url", "https://source.example.com", + "--source-api-key", "sk-source-123", + "--dry-run" + ]) + + assert result.exit_code == 0 + assert "Found 2 keys in source instance" in result.output + assert "DRY RUN MODE" in result.output + assert "test-key-1" in result.output + assert "user1@example.com" in result.output + assert "test-key-2" in result.output + assert "user2@example.com" in result.output + + # Verify source client was called (pagination stops early when fewer keys than page_size) + assert mock_source_instance.list.call_count >= 1 + mock_source_instance.list.assert_any_call(return_full_object=True, page=1, size=100) + + +def test_keys_import_actual_import_success(mock_keys_client, cli_runner): + """Test successful actual import of keys""" + # Create separate mock instances for source and destination + with patch("litellm.proxy.client.cli.commands.keys.KeysManagementClient") as MockClient: + mock_source_instance = MockClient.return_value + mock_dest_instance = MockClient.return_value + + # Configure source client + mock_source_instance.list.side_effect = [ + { + "keys": [ + { + "key_alias": "import-key-1", + "user_id": "user1@example.com", + "models": ["gpt-4"], + "spend": 100.0, + "team_id": "team-1" + } + ] + }, + {"keys": []} # Empty second page + ] + + # Configure destination client + mock_dest_instance.generate.return_value = { + "key": "sk-new-generated-key", + "status": "success" + } + + result = cli_runner.invoke(cli, [ + "keys", "import", + "--source-base-url", "https://source.example.com", + "--source-api-key", "sk-source-123" + ]) + + assert result.exit_code == 0 + assert "Found 1 keys in source instance" in result.output + assert "✓ Imported key: import-key-1" in result.output + assert "Successfully imported: 1" in result.output + assert "Failed to import: 0" in result.output + + # Verify generate was called with correct parameters + mock_dest_instance.generate.assert_called_once_with( + models=["gpt-4"], + spend=100.0, + key_alias="import-key-1", + team_id="team-1", + user_id="user1@example.com" + ) + + +def test_keys_import_pagination_handling(mock_keys_client, cli_runner): + """Test that import correctly handles pagination to get all keys""" + mock_source_instance = mock_keys_client.return_value + mock_source_instance.list.side_effect = [ + {"keys": [{"key_alias": f"key-{i}", "user_id": f"user{i}@example.com"} for i in range(100)]}, # Page 1: 100 keys + {"keys": [{"key_alias": f"key-{i}", "user_id": f"user{i}@example.com"} for i in range(100, 150)]}, # Page 2: 50 keys + {"keys": []} # Page 3: Empty + ] + + result = cli_runner.invoke(cli, [ + "keys", "import", + "--source-base-url", "https://source.example.com", + "--dry-run" + ]) + + assert result.exit_code == 0 + assert "Fetched page 1: 100 keys" in result.output + assert "Fetched page 2: 50 keys" in result.output + assert "Found 150 keys in source instance" in result.output + + # Verify pagination calls (stops early when fewer keys than page_size) + assert mock_source_instance.list.call_count >= 2 + mock_source_instance.list.assert_any_call(return_full_object=True, page=1, size=100) + mock_source_instance.list.assert_any_call(return_full_object=True, page=2, size=100) + + +def test_keys_import_created_since_filter(mock_keys_client, cli_runner): + """Test that --created-since filter works correctly""" + mock_source_instance = mock_keys_client.return_value + mock_source_instance.list.side_effect = [ + { + "keys": [ + { + "key_alias": "old-key", + "user_id": "user1@example.com", + "created_at": "2024-01-01T10:00:00Z" # Before filter + }, + { + "key_alias": "new-key", + "user_id": "user2@example.com", + "created_at": "2024-07-08T10:00:00Z" # After filter + } + ] + }, + {"keys": []} + ] + + result = cli_runner.invoke(cli, [ + "keys", "import", + "--source-base-url", "https://source.example.com", + "--created-since", "2024-07-07_18:19", + "--dry-run" + ]) + + assert result.exit_code == 0 + assert "Filtered 2 keys to 1 keys created since 2024-07-07_18:19" in result.output + assert "Found 1 keys in source instance" in result.output + assert "new-key" in result.output + assert "old-key" not in result.output + + +def test_keys_import_created_since_date_only_format(mock_keys_client, cli_runner): + """Test --created-since with date-only format (YYYY-MM-DD)""" + mock_source_instance = mock_keys_client.return_value + mock_source_instance.list.side_effect = [ + { + "keys": [ + { + "key_alias": "test-key", + "user_id": "user@example.com", + "created_at": "2024-07-08T10:00:00Z" + } + ] + }, + {"keys": []} + ] + + result = cli_runner.invoke(cli, [ + "keys", "import", + "--source-base-url", "https://source.example.com", + "--created-since", "2024-07-07", # Date only format + "--dry-run" + ]) + + assert result.exit_code == 0 + assert "Filtered 1 keys to 1 keys created since 2024-07-07" in result.output + + +def test_keys_import_no_keys_found(mock_keys_client, cli_runner): + """Test handling when no keys are found in source instance""" + mock_source_instance = mock_keys_client.return_value + mock_source_instance.list.return_value = {"keys": []} + + result = cli_runner.invoke(cli, [ + "keys", "import", + "--source-base-url", "https://source.example.com", + "--dry-run" + ]) + + assert result.exit_code == 0 + assert "No keys found in source instance" in result.output + + +def test_keys_import_invalid_date_format(cli_runner): + """Test error handling for invalid --created-since date format""" + result = cli_runner.invoke(cli, [ + "keys", "import", + "--source-base-url", "https://source.example.com", + "--created-since", "invalid-date", + "--dry-run" + ]) + + assert result.exit_code != 0 + assert "Invalid date format" in result.output + assert "Use YYYY-MM-DD_HH:MM or YYYY-MM-DD" in result.output + + +def test_keys_import_source_api_error(mock_keys_client, cli_runner): + """Test error handling when source API returns an error""" + mock_source_instance = mock_keys_client.return_value + mock_source_instance.list.side_effect = Exception("Source API Error") + + result = cli_runner.invoke(cli, [ + "keys", "import", + "--source-base-url", "https://source.example.com", + "--dry-run" + ]) + + assert result.exit_code != 0 + assert "Source API Error" in result.output + + +def test_keys_import_partial_failure(mock_keys_client, cli_runner): + """Test handling when some keys fail to import""" + with patch("litellm.proxy.client.cli.commands.keys.KeysManagementClient") as MockClient: + mock_source_instance = MockClient.return_value + mock_dest_instance = MockClient.return_value + + # Source returns 2 keys + mock_source_instance.list.side_effect = [ + { + "keys": [ + {"key_alias": "success-key", "user_id": "user1@example.com"}, + {"key_alias": "fail-key", "user_id": "user2@example.com"} + ] + }, + {"keys": []} + ] + + # Destination: first succeeds, second fails + mock_dest_instance.generate.side_effect = [ + {"key": "sk-new-key", "status": "success"}, + Exception("Import failed for this key") + ] + + result = cli_runner.invoke(cli, [ + "keys", "import", + "--source-base-url", "https://source.example.com" + ]) + + assert result.exit_code == 0 # Command completes even with partial failures + assert "✓ Imported key: success-key" in result.output + assert "✗ Failed to import key fail-key" in result.output + assert "Successfully imported: 1" in result.output + assert "Failed to import: 1" in result.output + assert "Total keys processed: 2" in result.output + + +def test_keys_import_missing_required_source_url(cli_runner): + """Test error when required --source-base-url is missing""" + result = cli_runner.invoke(cli, [ + "keys", "import", + "--dry-run" + ]) + + assert result.exit_code != 0 + assert "Missing option" in result.output or "required" in result.output.lower() + + +def test_keys_import_with_all_key_properties(mock_keys_client, cli_runner): + """Test import preserves all key properties (models, aliases, config, etc.)""" + with patch("litellm.proxy.client.cli.commands.keys.KeysManagementClient") as MockClient: + mock_source_instance = MockClient.return_value + mock_dest_instance = MockClient.return_value + + mock_source_instance.list.side_effect = [ + { + "keys": [ + { + "key_alias": "full-key", + "user_id": "user@example.com", + "team_id": "team-123", + "budget_id": "budget-456", + "models": ["gpt-4", "gpt-3.5-turbo"], + "aliases": {"custom-model": "gpt-4"}, + "spend": 50.0, + "config": {"max_tokens": 1000} + } + ] + }, + {"keys": []} + ] + + mock_dest_instance.generate.return_value = {"key": "sk-imported", "status": "success"} + + result = cli_runner.invoke(cli, [ + "keys", "import", + "--source-base-url", "https://source.example.com" + ]) + + assert result.exit_code == 0 + + # Verify all properties were passed to generate + mock_dest_instance.generate.assert_called_once_with( + models=["gpt-4", "gpt-3.5-turbo"], + aliases={"custom-model": "gpt-4"}, + spend=50.0, + key_alias="full-key", + team_id="team-123", + user_id="user@example.com", + budget_id="budget-456", + config={"max_tokens": 1000} + )