Change of policy: keys.json over-rides environment variables, closes #158

docs/setup.md

@@ -101,11 +101,12 @@ Keys can also be set using an environment variable. These are different for diff
 
 For OpenAI models the key will be read from the `OPENAI_API_KEY` environment variable.
 
-The environment variable will be used only if no `--key` option is passed to the command.
+The environment variable will be used if no `--key` option is passed to the command and there is not a key configured in `keys.json`
 
-If no environment variable is found, the tool will fall back to checking `keys.json`.
-
-You can force the tool to use the key from `keys.json` even if an environment variable has also been set using `llm "prompt" --key openai`.
+To use an environment variable in place of the `keys.json` key run the prompt like this:
+```bash
+llm 'my prompt' --key $OPENAI_API_KEY
+```
 
 ## Configuration
 

llm/__init__.py

@@ -111,11 +111,14 @@ def get_key(
     if explicit_key:
         # User specified a key that's not an alias, use that
         return explicit_key
-    # Environment variables over-ride the default key
+    # Stored key over-rides environment variables over-ride the default key
+    if key_alias in stored_keys:
+        return stored_keys[key_alias]
+    # Finally try environment variable
     if env_var and os.environ.get(env_var):
         return os.environ[env_var]
-    # Return the key stored for the default alias
-    return stored_keys.get(key_alias)
+    # Couldn't find it
+    return None
 
 
 def load_keys():

llm/models.py

@@ -5,7 +5,6 @@ import re
 import time
 from typing import Any, Dict, Iterator, List, Optional, Set
 from abc import ABC, abstractmethod
-import os
 import json
 from pydantic import BaseModel
 from ulid import ULID
@@ -220,15 +219,24 @@ class Model(ABC):
         pass
 
     def get_key(self):
-        if self.needs_key is None:
-            return None
-        if self.key is not None:
-            return self.key
-        if self.key_env_var is not None:
-            key = os.environ.get(self.key_env_var)
-            if key:
-                return key
+        from llm import get_key
 
+        if self.needs_key is None:
+            # This model doesn't use an API key
+            return None
+
+        if self.key is not None:
+            # Someone already set model.key='...'
+            return self.key
+
+        # Attempt to load a key using llm.get_key()
+        key = get_key(
+            explicit_key=None, key_alias=self.needs_key, env_var=self.key_env_var
+        )
+        if key:
+            return key
+
+        # Show a useful error message
         message = "No key found - add one using 'llm keys set {}'".format(
             self.needs_key
         )

tests/test_keys.py

@@ -48,15 +48,11 @@ def test_uses_correct_key(mocked_openai, monkeypatch, tmpdir):
     user_dir = tmpdir / "user-dir"
     pathlib.Path(user_dir).mkdir()
     keys_path = user_dir / "keys.json"
-    keys_path.write_text(
-        json.dumps(
-            {
-                "openai": "from-keys-file",
-                "other": "other-key",
-            }
-        ),
-        "utf-8",
-    )
+    KEYS = {
+        "openai": "from-keys-file",
+        "other": "other-key",
+    }
+    keys_path.write_text(json.dumps(KEYS), "utf-8")
     monkeypatch.setenv("LLM_USER_PATH", str(user_dir))
     monkeypatch.setenv("OPENAI_API_KEY", "from-env")
 
@@ -66,21 +62,26 @@ def test_uses_correct_key(mocked_openai, monkeypatch, tmpdir):
         ] == "Bearer {}".format(key)
 
     runner = CliRunner()
-    # Called without --key uses environment variable
+
+    # Called without --key uses stored key
     result = runner.invoke(cli, ["hello", "--no-stream"], catch_exceptions=False)
     assert result.exit_code == 0
-    assert_key("from-env")
-    # Called without --key and with no environment variable uses keys.json
-    monkeypatch.setenv("OPENAI_API_KEY", "")
+    assert_key("from-keys-file")
+
+    # Called without --key and without keys.json uses environment variable
+    keys_path.write_text("{}", "utf-8")
     result2 = runner.invoke(cli, ["hello", "--no-stream"], catch_exceptions=False)
     assert result2.exit_code == 0
-    assert_key("from-keys-file")
+    assert_key("from-env")
+    keys_path.write_text(json.dumps(KEYS), "utf-8")
+
     # Called with --key name-in-keys.json uses that value
     result3 = runner.invoke(
         cli, ["hello", "--key", "other", "--no-stream"], catch_exceptions=False
     )
     assert result3.exit_code == 0
     assert_key("other-key")
+
     # Called with --key something-else uses exactly that
     result4 = runner.invoke(
         cli, ["hello", "--key", "custom-key", "--no-stream"], catch_exceptions=False