scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

Add models import command (#10581)

Browse source 
* Add --only-models-matching-regex option

to `models import` which only processes models where
`litelllm_params.model` matches the regex

* Add test_models_import_only_models_matching_regex

* Print each model we're importing

* Add --only-access-groups-matching-regex option

to `models import` which only processes models where at least one item
in `model_info.access_groups` matches the regex. Add a unit test.

* Add `models import` examples to README.md

Add `models import` examples to proxy/client/cli/README.md

* ruff format litellm/proxy/client/cli/commands/models.py

* Make `models import` display tabular output

* models import refactoring

* Fix failing tests in test_models_commands.py

* Refactor import_models to make it shorter and more readable

* Extract from `import_models` a function called `get_model_list_from_yaml_file`

* Fix mypy error

* Add more specific typing

for better understandability and Intellisense

* More import_models refactoring

* More refactoring

* More refactoring

* Write unit tests for format_iso_datetime_str

* Add more unit tests

* ruff format tests/litellm/proxy/client/cli/test_models_commands.py

* ruff format litellm/proxy/client/cli/commands/models.py

* Make test_format_timestamp use UTC time
This commit is contained in:
Marc Abramowitz 2025-05-14 23:05:07 -07:00 • committed by GitHub
parent b88e56ebde
commit 322b67833b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 439 additions and 79 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

59
litellm/proxy/client/cli/README.md
View file

@ -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 <regex>`: Only import models where `litellm_params.model` matches the given regex.
- `--only-access-groups-matching-regex <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:

161
litellm/proxy/client/cli/commands/models.py
View file

@ -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)

298
tests/litellm/proxy/client/cli/test_models_commands.py
View file

@ -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()
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