diff --git a/docs/aliases.md b/docs/aliases.md index 3fa5c53..1fe531b 100644 --- a/docs/aliases.md +++ b/docs/aliases.md @@ -53,11 +53,8 @@ Example output: { "3.5": "gpt-3.5-turbo", "chatgpt": "gpt-3.5-turbo", - "chatgpt-16k": "gpt-3.5-turbo-16k", - "3.5-16k": "gpt-3.5-turbo-16k", "4": "gpt-4", "gpt4": "gpt-4", - "4-32k": "gpt-4-32k", "ada": "ada-002" } ``` @@ -67,12 +64,15 @@ Example output: The `llm aliases set ` command can be used to add a new alias: ```bash -llm aliases set turbo gpt-3.5-turbo-16k +llm aliases set mini gpt-4o-mini ``` -Now you can run the `gpt-3.5-turbo-16k` model using the `turbo` alias like this: - +You can also pass one or more `-q search` options to set an alias on the first model matching those search terms: ```bash -llm -m turbo 'An epic Greek-style saga about a cheesecake that builds a SQL database from scratch' +llm aliases set mini -q 4o -q mini +``` +Now you can run the `gpt-4o-mini` model using the `mini` alias like this: +```bash +llm -m mini 'An epic Greek-style saga about a cheesecake that builds a SQL database from scratch' ``` Aliases can be set for both regular models and {ref}`embedding models ` using the same command. To set an alias of `oai` for the OpenAI `ada-002` embedding model use this: ```bash @@ -92,7 +92,7 @@ Output: The `llm aliases remove ` command will remove the specified alias: ```bash -llm aliases remove turbo +llm aliases remove mini ``` ## Viewing the aliases file diff --git a/docs/help.md b/docs/help.md index b3614ee..b9d6664 100644 --- a/docs/help.md +++ b/docs/help.md @@ -444,16 +444,22 @@ Options: (help-aliases-set)= #### llm aliases set --help ``` -Usage: llm aliases set [OPTIONS] ALIAS MODEL_ID +Usage: llm aliases set [OPTIONS] ALIAS [MODEL_ID] Set an alias for a model Example usage: - $ llm aliases set turbo gpt-3.5-turbo + llm aliases set mini gpt-4o-mini + + Alternatively you can omit the model ID and specify one or more -q options. + The first model matching all of those query strings will be used. + + llm aliases set mini -q 4o -q mini Options: - --help Show this message and exit. + -q, --query TEXT Set alias for model matching these strings + --help Show this message and exit. ``` (help-aliases-remove)= diff --git a/llm/cli.py b/llm/cli.py index 09c9bee..01e6e6b 100644 --- a/llm/cli.py +++ b/llm/cli.py @@ -1300,17 +1300,51 @@ def aliases_list(json_): @aliases.command(name="set") @click.argument("alias") -@click.argument("model_id") -def aliases_set(alias, model_id): +@click.argument("model_id", required=False) +@click.option( + "-q", + "--query", + multiple=True, + help="Set alias for model matching these strings", +) +def aliases_set(alias, model_id, query): """ Set an alias for a model Example usage: \b - $ llm aliases set turbo gpt-3.5-turbo + llm aliases set mini gpt-4o-mini + + Alternatively you can omit the model ID and specify one or more -q options. + The first model matching all of those query strings will be used. + + \b + llm aliases set mini -q 4o -q mini """ - set_alias(alias, model_id) + if not model_id: + if not query: + raise click.ClickException( + "You must provide a model_id or at least one -q option" + ) + # Search for the first model matching all query strings + found = None + for model_with_aliases in get_models_with_aliases(): + if all(model_with_aliases.matches(q) for q in query): + found = model_with_aliases + break + if not found: + raise click.ClickException( + "No model found matching query: " + ", ".join(query) + ) + model_id = found.model.model_id + set_alias(alias, model_id) + click.echo( + f"Alias '{alias}' set to model '{model_id}'", + err=True, + ) + else: + set_alias(alias, model_id) @aliases.command(name="remove") diff --git a/tests/test_aliases.py b/tests/test_aliases.py index b08871a..853810b 100644 --- a/tests/test_aliases.py +++ b/tests/test_aliases.py @@ -71,14 +71,26 @@ def test_cli_aliases_list_json(args): ) -def test_cli_aliases_set(user_path): +@pytest.mark.parametrize( + "args,expected,expected_error", + ( + (["foo", "bar"], {"foo": "bar"}, None), + (["foo", "-q", "mo"], {"foo": "mock"}, None), + (["foo", "-q", "mog"], None, "No model found matching query: mog"), + ), +) +def test_cli_aliases_set(user_path, args, expected, expected_error): # Should be not aliases.json at start assert not (user_path / "aliases.json").exists() runner = CliRunner() - result = runner.invoke(cli, ["aliases", "set", "foo", "bar"]) - assert result.exit_code == 0 - assert (user_path / "aliases.json").exists() - assert json.loads((user_path / "aliases.json").read_text("utf-8")) == {"foo": "bar"} + result = runner.invoke(cli, ["aliases", "set"] + args) + if not expected_error: + assert result.exit_code == 0 + assert (user_path / "aliases.json").exists() + assert json.loads((user_path / "aliases.json").read_text("utf-8")) == expected + else: + assert result.exit_code == 1 + assert result.output.strip() == f"Error: {expected_error}" def test_cli_aliases_path(user_path):