From 2e90a30c4ff6fdd7a0c140d0aadc8b7aa3c5bb85 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 3 Sep 2023 17:55:56 -0700
Subject: [PATCH] Add embedding model aliases to llm aliases, refs #192

---
 docs/aliases.md       |  1 +
 llm/cli.py            | 22 ++++++++++++++++------
 tests/test_aliases.py |  6 ++++++
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/docs/aliases.md b/docs/aliases.md
index 843ac5b..4f35177 100644
--- a/docs/aliases.md
+++ b/docs/aliases.md
@@ -26,6 +26,7 @@ chatgpt-16k : gpt-3.5-turbo-16k
 4           : gpt-4
 gpt4        : gpt-4
 4-32k       : gpt-4-32k
+ada         : ada-002 (embedding)
 ```
 <!-- [[[end]]] -->
 
diff --git a/llm/cli.py b/llm/cli.py
index d985fe8..d0f4808 100644
--- a/llm/cli.py
+++ b/llm/cli.py
@@ -11,6 +11,7 @@ from llm import (
     UnknownModelError,
     encode,
     get_embedding_models_with_aliases,
+    get_embedding_model_aliases,
     get_embedding_model,
     get_key,
     get_plugins,
@@ -735,14 +736,23 @@ def aliases_list(json_):
     to_output = []
     for alias, model in get_model_aliases().items():
         if alias != model.model_id:
-            to_output.append((alias, model.model_id))
+            to_output.append((alias, model.model_id, ""))
+    for alias, embedding_model in get_embedding_model_aliases().items():
+        if alias != embedding_model.model_id:
+            to_output.append((alias, embedding_model.model_id, "embedding"))
     if json_:
-        click.echo(json.dumps({key: value for key, value in to_output}, indent=4))
+        click.echo(
+            json.dumps({key: value for key, value, type_ in to_output}, indent=4)
+        )
         return
-    max_alias_length = max(len(a) for a, _ in to_output)
-    fmt = "{alias:<" + str(max_alias_length) + "} : {model_id}"
-    for alias, model_id in to_output:
-        click.echo(fmt.format(alias=alias, model_id=model_id))
+    max_alias_length = max(len(a) for a, _, _ in to_output)
+    fmt = "{alias:<" + str(max_alias_length) + "} : {model_id}{type_}"
+    for alias, model_id, type_ in to_output:
+        click.echo(
+            fmt.format(
+                alias=alias, model_id=model_id, type_=f" ({type_})" if type_ else ""
+            )
+        )
 
 
 @aliases.command(name="set")
diff --git a/tests/test_aliases.py b/tests/test_aliases.py
index bd50a91..4647b46 100644
--- a/tests/test_aliases.py
+++ b/tests/test_aliases.py
@@ -25,6 +25,7 @@ def test_remove_alias():
 
 @pytest.mark.parametrize("args", (["aliases", "list"], ["aliases"]))
 def test_cli_aliases_list(args):
+    llm.set_alias("e-demo", "embed-demo")
     runner = CliRunner()
     result = runner.invoke(cli, args)
     assert result.exit_code == 0
@@ -36,11 +37,14 @@ def test_cli_aliases_list(args):
         "4           : gpt-4\n"
         "gpt4        : gpt-4\n"
         "4-32k       : gpt-4-32k\n"
+        "e-demo      : embed-demo (embedding)\n"
+        "ada         : ada-002 (embedding)\n"
     )
 
 
 @pytest.mark.parametrize("args", (["aliases", "list"], ["aliases"]))
 def test_cli_aliases_list_json(args):
+    llm.set_alias("e-demo", "embed-demo")
     runner = CliRunner()
     result = runner.invoke(cli, args + ["--json"])
     assert result.exit_code == 0
@@ -52,6 +56,8 @@ def test_cli_aliases_list_json(args):
         "4": "gpt-4",
         "gpt4": "gpt-4",
         "4-32k": "gpt-4-32k",
+        "ada": "ada-002",
+        "e-demo": "embed-demo",
     }