diff --git a/docs/aliases.md b/docs/aliases.md index f724326..bf343e0 100644 --- a/docs/aliases.md +++ b/docs/aliases.md @@ -46,3 +46,16 @@ Example output: "4-32k": "gpt-4-32k" } ``` + +## Adding a new alias + +The `llm aliases set ` command can be used to add a new alias: + +```bash +llm aliases set turbo gpt-3.5-turbo-16k +``` +Now you can run the `gpt-3.5-turbo-16k` model using the `turbo` alias like this: + +```bash +llm -m turbo 'An epic Greek-style saga about a cheesecake that builds a SQL database from scratch' +``` diff --git a/llm/cli.py b/llm/cli.py index d6d740c..54cacd7 100644 --- a/llm/cli.py +++ b/llm/cli.py @@ -298,7 +298,7 @@ def keys_path_command(): @keys.command(name="set") @click.argument("name") @click.option("--value", prompt="Enter key", hide_input=True, help="Value to set") -def set_(name, value): +def keys_set(name, value): """ Save a key in the keys.json file @@ -576,6 +576,30 @@ def aliases_list(json_): click.echo(fmt.format(alias=alias, model_id=model_id)) +@aliases.command(name="set") +@click.argument("alias") +@click.argument("model_id") +def aliases_set(alias, model_id): + """ + Set an alias for a model + + Example usage: + + \b + $ llm aliases set turbo gpt-3.5-turbo + """ + path = user_dir() / "aliases.json" + path.parent.mkdir(parents=True, exist_ok=True) + if not path.exists(): + path.write_text("{}\n") + try: + current = json.loads(path.read_text()) + except json.decoder.JSONDecodeError as ex: + raise click.ClickException("aliases.json is invalid: {}".format(ex)) + current[alias] = model_id + path.write_text(json.dumps(current, indent=4) + "\n") + + @cli.command(name="plugins") def plugins_list(): "List installed plugins" diff --git a/tests/test_aliases.py b/tests/test_aliases.py index f9ff964..adbaa68 100644 --- a/tests/test_aliases.py +++ b/tests/test_aliases.py @@ -31,3 +31,13 @@ def test_aliases_list_json(): "gpt4": "gpt-4", "4-32k": "gpt-4-32k", } + + +def test_aliases_set(user_path): + # 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"}