Skip to content

Commit

Permalink
Change of policy: keys.json over-rides environment variables, closes #…
Browse files Browse the repository at this point in the history
  • Loading branch information
simonw committed Aug 21, 2023
1 parent 341dbce commit dff36f0
Show file tree
Hide file tree
Showing 4 changed files with 39 additions and 26 deletions.
9 changes: 5 additions & 4 deletions docs/setup.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
9 changes: 6 additions & 3 deletions llm/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down
18 changes: 13 additions & 5 deletions llm/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
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
Expand Down Expand Up @@ -220,15 +219,24 @@ class Options(_Options):
pass

def get_key(self):
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
if self.key_env_var is not None:
key = os.environ.get(self.key_env_var)
if key:
return 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
)
Expand Down
29 changes: 15 additions & 14 deletions tests/test_keys.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand All @@ -66,21 +62,26 @@ def assert_key(key):
] == "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
Expand Down

0 comments on commit dff36f0

Please sign in to comment.