llm aliasess set command, refs #151

docs/aliases.md

@@ -46,3 +46,16 @@ Example output:
     "4-32k": "gpt-4-32k"
 }
 ```
+
+## Adding a new alias
+
+The `llm aliases set <alias> <model-id>` 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'
+```

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"

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"}