From 53289e3767d2ca03e3f3919bc1e9b84a47bc61cd Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sat, 12 Aug 2023 09:05:02 -0700
Subject: [PATCH] llm aliases remove command, refs #151

---
 docs/aliases.md       |  9 +++++++++
 llm/cli.py            | 25 +++++++++++++++++++++++++
 tests/test_aliases.py |  8 ++++++++
 3 files changed, 42 insertions(+)

diff --git a/docs/aliases.md b/docs/aliases.md
index 196cd1d..fc74c95 100644
--- a/docs/aliases.md
+++ b/docs/aliases.md
@@ -60,6 +60,15 @@ Now you can run the `gpt-3.5-turbo-16k` model using the `turbo` alias like this:
 llm -m turbo 'An epic Greek-style saga about a cheesecake that builds a SQL database from scratch'
 ```
 
+## Removing an alias
+
+The `llm aliases remove <alias>` command will remove the specified alias:
+
+```bash
+llm aliases remove turbo
+```
+This can also be used for aliases that were introduced by other plugins as opposed to the `llm aliases set` command.
+
 ## Viewing the aliases file
 
 Aliases are stored in an `aliases.json` file in the LLM configuration directory.
diff --git a/llm/cli.py b/llm/cli.py
index 0b3af13..6ad2a8e 100644
--- a/llm/cli.py
+++ b/llm/cli.py
@@ -600,6 +600,31 @@ def aliases_set(alias, model_id):
     path.write_text(json.dumps(current, indent=4) + "\n")
 
 
+@aliases.command(name="remove")
+@click.argument("alias")
+def aliases_remove(alias):
+    """
+    Remove an alias
+
+    Example usage:
+
+    \b
+        $ llm aliases remove 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))
+    if alias not in current:
+        raise click.ClickException("Alias not found: {}".format(alias))
+    del current[alias]
+    path.write_text(json.dumps(current, indent=4) + "\n")
+
+
 @aliases.command(name="path")
 def aliases_path():
     "Output the path to the aliases.json file"
diff --git a/tests/test_aliases.py b/tests/test_aliases.py
index e051715..f03c560 100644
--- a/tests/test_aliases.py
+++ b/tests/test_aliases.py
@@ -48,3 +48,11 @@ def test_aliases_path(user_path):
     result = runner.invoke(cli, ["aliases", "path"])
     assert result.exit_code == 0
     assert result.output.strip() == str(user_path / "aliases.json")
+
+
+def test_aliases_remove(user_path):
+    (user_path / "aliases.json").write_text(json.dumps({"foo": "bar"}), "utf-8")
+    runner = CliRunner()
+    result = runner.invoke(cli, ["aliases", "remove", "foo"])
+    assert result.exit_code == 0
+    assert json.loads((user_path / "aliases.json").read_text("utf-8")) == {}