diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 77724245b29..880b0b39720 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -123,6 +123,52 @@ Options: - `--param`, `-p`: Model parameters in key=value format (can be specified multiple times) - `--info`, `-i`: Model info in key=value format (can be specified multiple times) +#### Import Models + +Import models from a YAML file: + +```bash +litellm-proxy models import models.yaml +``` + +Options: + +- `--dry-run`: Show what would be imported without making any changes. +- `--only-models-matching-regex `: Only import models where `litellm_params.model` matches the given regex. +- `--only-access-groups-matching-regex `: Only import models where at least one item in `model_info.access_groups` matches the given regex. + +Examples: + +1. Import all models from a YAML file: + +```bash +litellm-proxy models import models.yaml +``` + +2. Dry run (show what would be imported): + +```bash +litellm-proxy models import models.yaml --dry-run +``` + +3. Only import models where the model name contains 'gpt': + +```bash +litellm-proxy models import models.yaml --only-models-matching-regex gpt +``` + +4. Only import models with access group containing 'beta': + +```bash +litellm-proxy models import models.yaml --only-access-groups-matching-regex beta +``` + +5. Combine both filters: + +```bash +litellm-proxy models import models.yaml --only-models-matching-regex gpt --only-access-groups-matching-regex beta +``` + ### Credentials Management The CLI provides commands for managing credentials on your LiteLLM proxy server: @@ -463,6 +509,19 @@ litellm-proxy users create --email a@b.com --role internal_user --alias "Alice" litellm-proxy users delete u1 u2 ``` +9. Import models from a YAML file (with filters): + +```bash +# Only import models where the model name contains 'gpt' +litellm-proxy models import models.yaml --only-models-matching-regex gpt + +# Only import models with access group containing 'beta' +litellm-proxy models import models.yaml --only-access-groups-matching-regex beta + +# Combine both filters +litellm-proxy models import models.yaml --only-models-matching-regex gpt --only-access-groups-matching-regex beta +``` + ## Error Handling The CLI will display appropriate error messages when: diff --git a/litellm/proxy/client/cli/commands/models.py b/litellm/proxy/client/cli/commands/models.py index 587eff993d7..4ff59e6be82 100644 --- a/litellm/proxy/client/cli/commands/models.py +++ b/litellm/proxy/client/cli/commands/models.py @@ -1,6 +1,10 @@ # stdlib imports -from typing import Optional, Literal, Any from datetime import datetime +import re +from typing import Optional, Literal, Any +import yaml +from dataclasses import dataclass +from collections import defaultdict # third party imports import click @@ -10,6 +14,38 @@ import rich from ... import Client +@dataclass +class ModelYamlInfo: + model_name: str + model_params: dict[str, Any] + model_info: dict[str, Any] + model_id: str + access_groups: list[str] + provider: str + + @property + def access_groups_str(self) -> str: + return ", ".join(self.access_groups) if self.access_groups else "" + + +def _get_model_info_obj_from_yaml(model: dict[str, Any]) -> ModelYamlInfo: + """Extract model info from a model dict and return as ModelYamlInfo dataclass.""" + model_name: str = model["model_name"] + model_params: dict[str, Any] = model["litellm_params"] + model_info: dict[str, Any] = model.get("model_info", {}) + model_id: str = model_params["model"] + access_groups = model_info.get("access_groups", []) + provider = model_id.split("/", 1)[0] if "/" in model_id else model_id + return ModelYamlInfo( + model_name=model_name, + model_params=model_params, + model_info=model_info, + model_id=model_id, + access_groups=access_groups, + provider=provider, + ) + + def format_iso_datetime_str(iso_datetime_str: Optional[str]) -> str: """Format an ISO format datetime string to human-readable date with minute resolution.""" if not iso_datetime_str: @@ -275,3 +311,126 @@ def update_model(ctx: click.Context, model_id: str, param: tuple[str, ...], info model_info=model_info, ) rich.print_json(data=result) + + +def _filter_model(model, model_regex, access_group_regex): + model_name = model.get("model_name") + model_params = model.get("litellm_params") + model_info = model.get("model_info", {}) + if not model_name or not model_params: + return False + model_id = model_params.get("model") + if not model_id or not isinstance(model_id, str): + return False + if model_regex and not model_regex.search(model_id): + return False + access_groups = model_info.get("access_groups", []) + if access_group_regex: + if not isinstance(access_groups, list): + return False + if not any(isinstance(group, str) and access_group_regex.search(group) for group in access_groups): + return False + return True + + +def _print_models_table(added_models: list[ModelYamlInfo], table_title: str): + if not added_models: + return + table = rich.table.Table(title=table_title) + table.add_column("Model Name", style="cyan") + table.add_column("Upstream Model", style="green") + table.add_column("Access Groups", style="magenta") + for m in added_models: + table.add_row(m.model_name, m.model_id, m.access_groups_str) + rich.print(table) + + +def _print_summary_table(provider_counts): + summary_table = rich.table.Table(title="Model Import Summary") + summary_table.add_column("Provider", style="cyan") + summary_table.add_column("Count", style="green") + + for provider, count in provider_counts.items(): + summary_table.add_row(str(provider), str(count)) + + total = sum(provider_counts.values()) + summary_table.add_row("[bold]Total[/bold]", f"[bold]{total}[/bold]") + + rich.print(summary_table) + + +def get_model_list_from_yaml_file(yaml_file: str) -> list[dict[str, Any]]: + """Load and validate the model list from a YAML file.""" + with open(yaml_file, "r") as f: + data = yaml.safe_load(f) + if not data or "model_list" not in data: + raise click.ClickException("YAML file must contain a 'model_list' key with a list of models.") + model_list = data["model_list"] + if not isinstance(model_list, list): + raise click.ClickException("'model_list' must be a list of model definitions.") + return model_list + + +def _get_filtered_model_list(model_list, only_models_matching_regex, only_access_groups_matching_regex): + """Return a list of models that pass the filter criteria.""" + model_regex = re.compile(only_models_matching_regex) if only_models_matching_regex else None + access_group_regex = re.compile(only_access_groups_matching_regex) if only_access_groups_matching_regex else None + return [model for model in model_list if _filter_model(model, model_regex, access_group_regex)] + + +def _import_models_get_table_title(dry_run: bool) -> str: + if dry_run: + return "Models that would be imported if [yellow]--dry-run[/yellow] was not provided" + else: + return "Models Imported" + + +@models.command("import") +@click.argument("yaml_file", type=click.Path(exists=True, dir_okay=False, readable=True)) +@click.option("--dry-run", is_flag=True, help="Show what would be imported without making any changes.") +@click.option( + "--only-models-matching-regex", + default=None, + help="Only import models where litellm_params.model matches the given regex.", +) +@click.option( + "--only-access-groups-matching-regex", + default=None, + help="Only import models where at least one item in model_info.access_groups matches the given regex.", +) +@click.pass_context +def import_models( + ctx: click.Context, + yaml_file: str, + dry_run: bool, + only_models_matching_regex: Optional[str], + only_access_groups_matching_regex: Optional[str], +) -> None: + """Import models from a YAML file and add them to the proxy.""" + provider_counts: dict[str, int] = defaultdict(int) + added_models: list[ModelYamlInfo] = [] + model_list = get_model_list_from_yaml_file(yaml_file) + filtered_model_list = _get_filtered_model_list( + model_list, only_models_matching_regex, only_access_groups_matching_regex + ) + + if not dry_run: + client = create_client(ctx) + + for model in filtered_model_list: + model_info_obj = _get_model_info_obj_from_yaml(model) + if not dry_run: + try: + client.models.new( + model_name=model_info_obj.model_name, + model_params=model_info_obj.model_params, + model_info=model_info_obj.model_info, + ) + except Exception: + pass # For summary, ignore errors + added_models.append(model_info_obj) + provider_counts[model_info_obj.provider] += 1 + + table_title = _import_models_get_table_title(dry_run) + _print_models_table(added_models, table_title) + _print_summary_table(provider_counts) diff --git a/tests/litellm/proxy/client/cli/test_models_commands.py b/tests/litellm/proxy/client/cli/test_models_commands.py index 62b6c42c5de..f05a20ee53d 100644 --- a/tests/litellm/proxy/client/cli/test_models_commands.py +++ b/tests/litellm/proxy/client/cli/test_models_commands.py @@ -1,6 +1,7 @@ # stdlib imports import json import os +import time from unittest.mock import patch # third party imports @@ -9,7 +10,11 @@ import pytest # local imports from litellm.proxy.client.cli import cli -from litellm.proxy.client.cli.commands.models import format_timestamp +from litellm.proxy.client.cli.commands.models import ( + format_timestamp, + format_iso_datetime_str, + format_cost_per_1k_tokens, +) @pytest.fixture @@ -28,10 +33,13 @@ def cli_runner(): @pytest.fixture(autouse=True) def mock_env(): """Fixture to set up environment variables for all tests""" - with patch.dict(os.environ, { - "LITELLM_PROXY_URL": "http://localhost:4000", - "LITELLM_PROXY_API_KEY": "sk-test" - }): + with patch.dict( + os.environ, + { + "LITELLM_PROXY_URL": "http://localhost:4000", + "LITELLM_PROXY_API_KEY": "sk-test", + }, + ): yield @@ -39,20 +47,10 @@ def mock_env(): def mock_models_list(mock_client): """Fixture to set up common mocking pattern for models list tests""" mock_client.return_value.models.list.return_value = [ - { - "id": "model-123", - "object": "model", - "created": 1699848889, - "owned_by": "organization-123" - }, - { - "id": "model-456", - "object": "model", - "created": 1699848890, - "owned_by": "organization-456" - } + {"id": "model-123", "object": "model", "created": 1699848889, "owned_by": "organization-123"}, + {"id": "model-456", "object": "model", "created": 1699848890, "owned_by": "organization-456"}, ] - + mock_client.assert_not_called() # Ensure clean slate return mock_client @@ -63,41 +61,53 @@ def mock_models_info(mock_client): mock_client.return_value.models.info.return_value = [ { "model_name": "gpt-4", - "litellm_params": { - "model": "gpt-4", - "litellm_credential_name": "openai-1" - }, + "litellm_params": {"model": "gpt-4", "litellm_credential_name": "openai-1"}, "model_info": { "id": "model-123", "created_at": "2025-04-29T21:31:43.843000+00:00", "updated_at": "2025-04-29T21:31:43.843000+00:00", "input_cost_per_token": 0.00001, - "output_cost_per_token": 0.00002 - } + "output_cost_per_token": 0.00002, + }, } ] - + mock_client.assert_not_called() return mock_client +@pytest.fixture +def force_utc_tz(): + """Fixture to force UTC timezone for tests that depend on system TZ.""" + old_tz = os.environ.get("TZ") + os.environ["TZ"] = "UTC" + if hasattr(time, "tzset"): + time.tzset() + yield + # Restore previous TZ + if old_tz is not None: + os.environ["TZ"] = old_tz + else: + if "TZ" in os.environ: + del os.environ["TZ"] + if hasattr(time, "tzset"): + time.tzset() + + def test_models_list_json_format(mock_models_list, cli_runner): """Test the models list command with JSON output format""" # Run the command result = cli_runner.invoke(cli, ["models", "list", "--format", "json"]) - + # Check that the command succeeded assert result.exit_code == 0 - + # Parse the output and verify it matches our mock data output_data = json.loads(result.output) assert output_data == mock_models_list.return_value.models.list.return_value - + # Verify the client was called correctly - mock_models_list.assert_called_once_with( - base_url="http://localhost:4000", - api_key="sk-test" - ) + mock_models_list.assert_called_once_with(base_url="http://localhost:4000", api_key="sk-test") mock_models_list.return_value.models.list.assert_called_once() @@ -105,10 +115,10 @@ def test_models_list_table_format(mock_models_list, cli_runner): """Test the models list command with table output format""" # Run the command result = cli_runner.invoke(cli, ["models", "list"]) - + # Check that the command succeeded assert result.exit_code == 0 - + # Verify the output contains expected table elements assert "ID" in result.output assert "Object" in result.output @@ -117,52 +127,43 @@ def test_models_list_table_format(mock_models_list, cli_runner): assert "model-123" in result.output assert "organization-123" in result.output assert format_timestamp(1699848889) in result.output - + # Verify the client was called correctly - mock_models_list.assert_called_once_with( - base_url="http://localhost:4000", - api_key="sk-test" - ) + mock_models_list.assert_called_once_with(base_url="http://localhost:4000", api_key="sk-test") mock_models_list.return_value.models.list.assert_called_once() def test_models_list_with_base_url(mock_models_list, cli_runner): """Test the models list command with custom base URL overriding env var""" custom_base_url = "http://custom.server:8000" - + # Run the command with custom base URL - result = cli_runner.invoke(cli, [ - "--base-url", custom_base_url, - "models", "list" - ]) - + result = cli_runner.invoke(cli, ["--base-url", custom_base_url, "models", "list"]) + # Check that the command succeeded assert result.exit_code == 0 - + # Verify the client was created with the custom base URL (overriding env var) mock_models_list.assert_called_once_with( base_url=custom_base_url, - api_key="sk-test" # Should still use env var for API key + api_key="sk-test", # Should still use env var for API key ) def test_models_list_with_api_key(mock_models_list, cli_runner): """Test the models list command with API key overriding env var""" custom_api_key = "custom-test-key" - + # Run the command with custom API key - result = cli_runner.invoke(cli, [ - "--api-key", custom_api_key, - "models", "list" - ]) - + result = cli_runner.invoke(cli, ["--api-key", custom_api_key, "models", "list"]) + # Check that the command succeeded assert result.exit_code == 0 - + # Verify the client was created with the custom API key (overriding env var) mock_models_list.assert_called_once_with( base_url="http://localhost:4000", # Should still use env var for base URL - api_key=custom_api_key + api_key=custom_api_key, ) @@ -170,38 +171,32 @@ def test_models_list_error_handling(mock_client, cli_runner): """Test error handling in the models list command""" # Configure mock to raise an exception mock_client.return_value.models.list.side_effect = Exception("API Error") - + # Run the command result = cli_runner.invoke(cli, ["models", "list"]) - + # Check that the command failed assert result.exit_code != 0 assert "API Error" in str(result.exception) - + # Verify the client was created with env var values - mock_client.assert_called_once_with( - base_url="http://localhost:4000", - api_key="sk-test" - ) + mock_client.assert_called_once_with(base_url="http://localhost:4000", api_key="sk-test") def test_models_info_json_format(mock_models_info, cli_runner): """Test the models info command with JSON output format""" # Run the command result = cli_runner.invoke(cli, ["models", "info", "--format", "json"]) - + # Check that the command succeeded assert result.exit_code == 0 - + # Parse the output and verify it matches our mock data output_data = json.loads(result.output) assert output_data == mock_models_info.return_value.models.info.return_value - + # Verify the client was called correctly with env var values - mock_models_info.assert_called_once_with( - base_url="http://localhost:4000", - api_key="sk-test" - ) + mock_models_info.assert_called_once_with(base_url="http://localhost:4000", api_key="sk-test") mock_models_info.return_value.models.info.assert_called_once() @@ -209,24 +204,171 @@ def test_models_info_table_format(mock_models_info, cli_runner): """Test the models info command with table output format""" # Run the command with default columns result = cli_runner.invoke(cli, ["models", "info"]) - + # Check that the command succeeded assert result.exit_code == 0 - + # Verify the output contains expected table elements assert "Public Model" in result.output assert "Upstream Model" in result.output assert "Updated At" in result.output assert "gpt-4" in result.output assert "2025-04-29 21:31" in result.output - + # Verify seconds and microseconds are not shown assert "21:31:43" not in result.output assert "843000" not in result.output - + # Verify the client was called correctly with env var values - mock_models_info.assert_called_once_with( - base_url="http://localhost:4000", - api_key="sk-test" - ) - mock_models_info.return_value.models.info.assert_called_once() \ No newline at end of file + mock_models_info.assert_called_once_with(base_url="http://localhost:4000", api_key="sk-test") + mock_models_info.return_value.models.info.assert_called_once() + + +def test_models_import_only_models_matching_regex(tmp_path, mock_client, cli_runner): + """Test the --only-models-matching-regex option for models import command""" + # Prepare a YAML file with a mix of models + yaml_content = { + "model_list": [ + {"model_name": "gpt-4-model", "litellm_params": {"model": "gpt-4"}, "model_info": {"id": "id-1"}}, + {"model_name": "gpt-3.5-model", "litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "id-2"}}, + {"model_name": "llama2-model", "litellm_params": {"model": "llama2"}, "model_info": {"id": "id-3"}}, + {"model_name": "other-model", "litellm_params": {"model": "other"}, "model_info": {"id": "id-4"}}, + ] + } + import yaml as pyyaml + + yaml_file = tmp_path / "models.yaml" + with open(yaml_file, "w") as f: + pyyaml.safe_dump(yaml_content, f) + + # Patch client.models.new to track calls + mock_new = mock_client.return_value.models.new + + # Only match models containing 'gpt' in their litellm_params.model + result = cli_runner.invoke(cli, ["models", "import", str(yaml_file), "--only-models-matching-regex", "gpt"]) + + # Should succeed + assert result.exit_code == 0 + # Only the two gpt models should be imported + calls = [call.kwargs["model_params"]["model"] for call in mock_new.call_args_list] + assert set(calls) == {"gpt-4", "gpt-3.5-turbo"} + # Should not include llama2 or other + assert "llama2" not in calls + assert "other" not in calls + # Output summary should mention the correct providers + assert "gpt-4".split("-")[0] in result.output or "gpt" in result.output + + +def test_models_import_only_access_groups_matching_regex(tmp_path, mock_client, cli_runner): + """Test the --only-access-groups-matching-regex option for models import command""" + # Prepare a YAML file with a mix of models + yaml_content = { + "model_list": [ + { + "model_name": "gpt-4-model", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "id-1", "access_groups": ["beta-models", "prod-models"]}, + }, + { + "model_name": "gpt-3.5-model", + "litellm_params": {"model": "gpt-3.5-turbo"}, + "model_info": {"id": "id-2", "access_groups": ["alpha-models"]}, + }, + { + "model_name": "llama2-model", + "litellm_params": {"model": "llama2"}, + "model_info": {"id": "id-3", "access_groups": ["beta-models"]}, + }, + { + "model_name": "other-model", + "litellm_params": {"model": "other"}, + "model_info": {"id": "id-4", "access_groups": ["other-group"]}, + }, + { + "model_name": "no-access-group-model", + "litellm_params": {"model": "no-access"}, + "model_info": {"id": "id-5"}, + }, + ] + } + import yaml as pyyaml + + yaml_file = tmp_path / "models.yaml" + with open(yaml_file, "w") as f: + pyyaml.safe_dump(yaml_content, f) + + # Patch client.models.new to track calls + mock_new = mock_client.return_value.models.new + + # Only match models with access_groups containing 'beta' + result = cli_runner.invoke(cli, ["models", "import", str(yaml_file), "--only-access-groups-matching-regex", "beta"]) + + # Should succeed + assert result.exit_code == 0 + # Only the two models with 'beta-models' in access_groups should be imported + calls = [call.kwargs["model_params"]["model"] for call in mock_new.call_args_list] + assert set(calls) == {"gpt-4", "llama2"} + # Should not include gpt-3.5, other, or no-access + assert "gpt-3.5-turbo" not in calls + assert "other" not in calls + assert "no-access" not in calls + # Output summary should mention the correct providers + assert "gpt-4".split("-")[0] in result.output or "gpt" in result.output + + +@pytest.mark.parametrize( + "input_str,expected", + [ + (None, ""), + ("", ""), + ("2024-05-01T12:34:56Z", "2024-05-01 12:34"), + ("2024-05-01T12:34:56+00:00", "2024-05-01 12:34"), + ("2024-05-01T12:34:56.123456+00:00", "2024-05-01 12:34"), + ("2024-05-01T12:34:56.123456Z", "2024-05-01 12:34"), + ("2024-05-01T12:34:56-04:00", "2024-05-01 12:34"), + ("2024-05-01", "2024-05-01 00:00"), + ("not-a-date", "not-a-date"), + ], +) +def test_format_iso_datetime_str(input_str, expected): + assert format_iso_datetime_str(input_str) == expected + + +@pytest.mark.parametrize( + "input_val,expected", + [ + (None, ""), + (1699848889, "2023-11-13 04:14"), + (1699848889.0, "2023-11-13 04:14"), + ("not-a-timestamp", "not-a-timestamp"), + ([1, 2, 3], "[1, 2, 3]"), + ], +) +def test_format_timestamp(input_val, expected, force_utc_tz): + actual = format_timestamp(input_val) + if actual != expected: + print(f"input: {input_val}, expected: {expected}, actual: {actual}") + assert actual == expected + + +@pytest.mark.parametrize( + "input_val,expected", + [ + (None, ""), + (0, "$0.0000"), + (0.0, "$0.0000"), + (0.00001, "$0.0100"), + (0.00002, "$0.0200"), + (1, "$1000.0000"), + (1.5, "$1500.0000"), + ("0.00001", "$0.0100"), + ("1.5", "$1500.0000"), + ("not-a-number", "not-a-number"), + (1e-10, "$0.0000"), + ], +) +def test_format_cost_per_1k_tokens(input_val, expected): + actual = format_cost_per_1k_tokens(input_val) + if actual != expected: + print(f"input: {input_val}, expected: {expected}, actual: {actual}") + assert actual == expected