feat: add --api-key flag to set provider API keys as env vars

This commit is contained in:
Paul Gauthier (aider) 2024-12-07 08:23:50 -08:00
parent 1d4918dfbf
commit 13ff038e58
3 changed files with 42 additions and 0 deletions

View file

@ -779,6 +779,13 @@ def get_parser(default_config_files, git_root):
help="Set an environment variable (can be used multiple times)",
default=[],
)
group.add_argument(
"--api-key",
action="append",
metavar="PROVIDER=KEY",
help="Set an API key for a provider (eg: --api-key anthropic=sk-123)",
default=[],
)
##########
group = parser.add_argument_group("Voice Settings")

View file

@ -537,6 +537,18 @@ def main(argv=None, input=None, output=None, force_git_root=None, return_coder=F
io.tool_output("Format should be: ENV_VAR_NAME=value")
return 1
# Process any API keys set via --api-key
if args.api_key:
for api_setting in args.api_key:
try:
provider, key = api_setting.split("=", 1)
env_var = f"{provider.strip().upper()}_API_KEY"
os.environ[env_var] = key.strip()
except ValueError:
io.tool_error(f"Invalid --api-key format: {api_setting}")
io.tool_output("Format should be: provider=key")
return 1
analytics = Analytics(logfile=args.analytics_log, permanently_disable=args.analytics_disable)
if args.analytics is not False:
if analytics.need_to_ask(args.analytics):

View file

@ -706,6 +706,29 @@ class TestMain(TestCase):
result = main(["--set-env", "INVALID_FORMAT", "--exit", "--yes"])
self.assertEqual(result, 1)
def test_api_key_single(self):
# Test setting a single API key
with GitTemporaryDirectory():
main(["--api-key", "anthropic=test-key", "--exit", "--yes"])
self.assertEqual(os.environ.get("ANTHROPIC_API_KEY"), "test-key")
def test_api_key_multiple(self):
# Test setting multiple API keys
with GitTemporaryDirectory():
main([
"--api-key", "anthropic=key1",
"--api-key", "openai=key2",
"--exit", "--yes"
])
self.assertEqual(os.environ.get("ANTHROPIC_API_KEY"), "key1")
self.assertEqual(os.environ.get("OPENAI_API_KEY"), "key2")
def test_api_key_invalid_format(self):
# Test invalid format handling
with GitTemporaryDirectory():
result = main(["--api-key", "INVALID_FORMAT", "--exit", "--yes"])
self.assertEqual(result, 1)
def test_invalid_edit_format(self):
with GitTemporaryDirectory():
with patch("aider.io.InputOutput.offer_url") as mock_offer_url: