From e959026051f4f5bafd8ac8967b6b62934cc97ddd Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 20 Aug 2023 23:11:47 -0700
Subject: [PATCH] llm keys command, closes #174

---
 docs/help.md       | 14 ++++++++++++--
 docs/setup.md      |  8 ++++++++
 llm/cli.py         | 19 ++++++++++++++++++-
 tests/test_keys.py | 12 ++++++++++++
 4 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/docs/help.md b/docs/help.md
index 2f33331b..ed7fb981 100644
--- a/docs/help.md
+++ b/docs/help.md
@@ -97,8 +97,18 @@ Options:
   --help  Show this message and exit.
 
 Commands:
-  path  Output the path to the keys.json file
-  set   Save a key in the keys.json file
+  list*  List names of all stored keys
+  path   Output the path to the keys.json file
+  set    Save a key in the keys.json file
+```
+#### llm keys list --help
+```
+Usage: llm keys list [OPTIONS]
+
+  List names of all stored keys
+
+Options:
+  --help  Show this message and exit.
 ```
 #### llm keys path --help
 ```
diff --git a/docs/setup.md b/docs/setup.md
index 8dda08b3..9ca845cc 100644
--- a/docs/setup.md
+++ b/docs/setup.md
@@ -63,11 +63,19 @@ Once stored, this key will be automatically used for subsequent calls to the API
 ```bash
 llm "Five ludicrous names for a pet lobster"
 ```
+
+You can list the names of keys that have been set using this command:
+
+```bash
+llm keys
+```
+
 Keys that are stored in this way live in a file called `keys.json`. This file is located at the path shown when you run the following command:
 
 ```bash
 llm keys path
 ```
+
 On macOS this will be `~/Library/Application Support/io.datasette.llm/keys.json`. On Linux it may be something like `~/.config/io.datasette.llm/keys.json`.
 
 ### Passing keys using the --key option
diff --git a/llm/cli.py b/llm/cli.py
index 8ed1e62b..ece01d12 100644
--- a/llm/cli.py
+++ b/llm/cli.py
@@ -293,11 +293,28 @@ def load_conversation(conversation_id: Optional[str]) -> Optional[Conversation]:
     return conversation
 
 
-@cli.group()
+@cli.group(
+    cls=DefaultGroup,
+    default="list",
+    default_if_no_args=True,
+)
 def keys():
     "Manage stored API keys for different models"
 
 
+@keys.command(name="list")
+def keys_list():
+    "List names of all stored keys"
+    path = user_dir() / "keys.json"
+    if not path.exists():
+        click.echo("No keys found")
+        return
+    keys = json.loads(path.read_text())
+    for key in sorted(keys.keys()):
+        if key != "// Note":
+            click.echo(key)
+
+
 @keys.command(name="path")
 def keys_path_command():
     "Output the path to the keys.json file"
diff --git a/tests/test_keys.py b/tests/test_keys.py
index a36fac70..46db22ff 100644
--- a/tests/test_keys.py
+++ b/tests/test_keys.py
@@ -32,6 +32,18 @@ def test_keys_set(monkeypatch, tmpdir):
     }
 
 
+@pytest.mark.parametrize("args", (["keys", "list"], ["keys"]))
+def test_keys_list(monkeypatch, tmpdir, args):
+    user_path = str(tmpdir / "user/keys")
+    monkeypatch.setenv("LLM_USER_PATH", user_path)
+    runner = CliRunner()
+    result = runner.invoke(cli, ["keys", "set", "openai"], input="foo")
+    assert result.exit_code == 0
+    result2 = runner.invoke(cli, args)
+    assert result2.exit_code == 0
+    assert result2.output.strip() == "openai"
+
+
 def test_uses_correct_key(mocked_openai, monkeypatch, tmpdir):
     user_dir = tmpdir / "user-dir"
     pathlib.Path(user_dir).mkdir()