From 458a59a8f4ba6aefac872c615bac901ac0176174 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sat, 12 Aug 2023 09:15:17 -0700
Subject: [PATCH] Actually register aliases from aliases.json, closes #151

---
 llm/__init__.py       | 14 +++++++++++++-
 tests/test_aliases.py | 10 ++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/llm/__init__.py b/llm/__init__.py
index 93c52995..6fa9d272 100644
--- a/llm/__init__.py
+++ b/llm/__init__.py
@@ -54,10 +54,22 @@ def get_plugins():
 def get_models_with_aliases() -> List["ModelWithAliases"]:
     model_aliases = []
 
+    # Include aliases from aliases.json
+    aliases_path = user_dir() / "aliases.json"
+    extra_model_aliases = {}
+    if aliases_path.exists():
+        configured_aliases = json.loads(aliases_path.read_text())
+        for alias, model_id in configured_aliases.items():
+            extra_model_aliases.setdefault(model_id, []).append(alias)
+
     def register(model, aliases=None):
-        model_aliases.append(ModelWithAliases(model, aliases or set()))
+        alias_list = list(aliases or [])
+        if model.model_id in extra_model_aliases:
+            alias_list.extend(extra_model_aliases[model.model_id])
+        model_aliases.append(ModelWithAliases(model, alias_list))
 
     pm.hook.register_models(register=register)
+
     return model_aliases
 
 
diff --git a/tests/test_aliases.py b/tests/test_aliases.py
index 4bf27a1d..8dc08a7e 100644
--- a/tests/test_aliases.py
+++ b/tests/test_aliases.py
@@ -64,3 +64,13 @@ def test_aliases_remove_invalid(user_path):
     result = runner.invoke(cli, ["aliases", "remove", "invalid"])
     assert result.exit_code == 1
     assert result.output == "Error: Alias not found: invalid\n"
+
+
+def test_aliases_are_registered(user_path):
+    (user_path / "aliases.json").write_text(
+        json.dumps({"foo": "bar", "turbo": "gpt-3.5-turbo"}), "utf-8"
+    )
+    runner = CliRunner()
+    result = runner.invoke(cli, ["models", "list"])
+    assert result.exit_code == 0
+    assert "gpt-3.5-turbo (aliases: 3.5, chatgpt, turbo)" in result.output