aikey

#!/usr/bin/env python3
"""
A minimal-dependency Python 3 script to manage and select AI API keys from environment variables,
then run a specified command with those keys set in its environment.

Goals and objectives (for a technical audience, concisely):
- Provide an easy way to choose which API key(s) to use at runtime.
- Support known providers like "openai" -> "OPENAI_API_KEY", "replicate" -> "REPLICATE_API_TOKEN", etc.
- Permits the user, via command line flags, to specify which keys (and/or providers) to look up.
- If multiple matching environment variables are found, the user can interactively select one.
- Once selected, the script sets those keys in the environment and executes a command.
- Verbosity is adjustable via repeated "-v" or "--verbose" flags.

Useful example:
    aikey -k openai,replicate -- echo "Testing environment..."

Future roadmap:
- Support .env files or other credential stores.
- Additional environment variable name transformations beyond prefix matching.
- More advanced interactive selection (e.g., filtering, searching, etc.).


DEMO:

```
$ # Here we use 'bash -c...' as simluation of program 
$ #    using environment values. Here we print them.

$ export OPENAI_API_KEY='aaa'
$ export OPENAI_API_KEY_BBB='bbb'
$ export OPENAI_API_KEY_CCC='ccc'
$ export ANTHROPIC_API_KEY='zoo'
$ export ANTHROPIC_API_KEY_AAA='foo'
$ export ANTHROPIC_API_KEY_BBB='bar'

$ # Here user asks `aikey` to "wrap" execution and select which enviornment values one wants to use. With --reveal-values flga, prints them to terminal.
$ aikey -k openai,anthropic --reveal-values bash -c 'echo "$OPENAI_API_KEY"; echo "$ANTHROPIC_API_KEY"'
1) OPENAI_API_KEY_BBB = bbb
2) OPENAI_API_KEY = aaa
3) OPENAI_API_KEY_CCC = ccc
Select a number, or 'q' to skip:
> 1
1) ANTHROPIC_API_KEY_AAA = foo
2) ANTHROPIC_API_KEY = zoo
3) ANTHROPIC_API_KEY_BBB = bar
Select a number, or 'q' to skip:
> 1
bbb
foo

$ aikey -k openai,anthropic --reveal-values bash -c 'echo "$OPENAI_API_KEY"; echo "$ANTHROPIC_API_KEY"'
1) OPENAI_API_KEY_BBB = bbb
2) OPENAI_API_KEY = aaa
3) OPENAI_API_KEY_CCC = ccc
Select a number, or 'q' to skip:
> 3
1) ANTHROPIC_API_KEY_AAA = foo
2) ANTHROPIC_API_KEY = zoo
3) ANTHROPIC_API_KEY_BBB = bar
Select a number, or 'q' to skip:
> 3
ccc
bar

$ # here example without --reveal-flag
$ aikey -k openai,anthropic bash -c 'echo "$OPENAI_API_KEY"; echo "$ANTHROPIC_API_KEY"'
1) OPENAI_API_KEY_BBB
2) OPENAI_API_KEY
3) OPENAI_API_KEY_CCC
Select a number, or 'q' to skip:
> 2
1) ANTHROPIC_API_KEY_AAA
2) ANTHROPIC_API_KEY
3) ANTHROPIC_API_KEY_BBB
Select a number, or 'q' to skip:
> 2
aaa
zoo
```
"""

import os
import sys
import argparse
import subprocess

KNOWN_PROVIDERS = {
    'openai': 'OPENAI_API_KEY',
    'replicate': 'REPLICATE_API_TOKEN',
    'anthropic': 'ANTHROPIC_API_KEY'
}

def parse_args():
    """
    Parse command line arguments.
    -k, --keys: comma-separated list of key names or known providers.
    -v, --verbose: increases verbosity level (each occurrence adds 1).
    --reveal-values: if set, display actual environment variable values in the selection list.
    -s, --shell-export: if set, output environment variable exports suitable for 'eval' in a shell.
    Everything after '--' or the first non-option argument is considered the command to run.
    """
    parser = argparse.ArgumentParser(
        description="Select AI API keys from environment and run a command with those keys set.",
        epilog="""
Additional usage notes for --shell-export:
------------------------------------------
If you want to 'activate' chosen keys in your current shell, you can define a
function in your .bashrc or .profile, for example:

    function eaikey() {
        eval "$(aikey --shell-export "$@" < /dev/tty 2>/dev/tty)"
    }

Then run:

    eaikey -k openai,replicate

and you'll be prompted (via stderr/stdin on your terminal) to choose the variables.
They will be set in your current shell session afterwards.
""",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        '-k', '--keys',
        help="Comma-separated list of keys or providers to load. Examples: openai,replicate,anthropic."
    )
    parser.add_argument(
        '-v', '--verbose', action='count', default=0,
        help="Increase verbosity level (use multiple times for more)."
    )
    parser.add_argument(
        '--reveal-values', action='store_true', default=False,
        help="If set, display actual environment variable values in the selection menu."
    )
    parser.add_argument(
        '-s', '--shell-export', action='store_true', default=False,
        help="If set, output shell export statements instead of running a command."
    )

    # We'll split out the command to run after the known flags.
    # The command can appear after a '--' or be all items that remain after known flags.
    args, remainder = parser.parse_known_args()

    # If remainder starts with '--', remove it.
    if len(remainder) > 0 and remainder[0] == '--':
        remainder = remainder[1:]

    args.command = remainder
    return args

def log_message(msg, level, current_verbosity, prefix="INFO"):
    """
    Print a message with a prefix to stderr if current_verbosity >= level.
    """
    if current_verbosity >= level:
        print(f"{prefix}: {msg}", file=sys.stderr)

def find_envvars_for_provider(provider_str):
    """
    Given a provider string (e.g. 'openai'), check if we have it in KNOWN_PROVIDERS.
    If so, retrieve the canonical environment variable name. If not, assume the user
    wants that exact string as the prefix.
    Then find environment variables that start with that prefix (case-sensitive).
    """
    if provider_str.lower() in KNOWN_PROVIDERS:
        prefix = KNOWN_PROVIDERS[provider_str.lower()]
    else:
        prefix = provider_str

    matches = [var for var in os.environ if var.startswith(prefix)]
    return matches, prefix

def select_envvar_from_list(matches, prefix, verbosity, reveal_values=False, shell_export=False):
    """
    If multiple matches are found, interactively let user pick one.
    If exactly one match is found, return it immediately.
    If none found, return None.

    Returns the name of the *selected* environment variable or None.
    """
    if not matches:
        log_message(f"No environment variable found matching prefix '{prefix}'.", 1, verbosity, "INFO")
        return None
    if len(matches) == 1:
        return matches[0]

    # Multiple found; prompt user
    log_message(f"Multiple environment variables match prefix '{prefix}':", 1, verbosity, "INFO")
    for i, m in enumerate(matches, start=1):
        if reveal_values:
            print(f"{i}) {m} = {os.environ[m]}", file=sys.stderr)
        else:
            print(f"{i}) {m}", file=sys.stderr)

    print("Select a number, or 'q' to skip:", file=sys.stderr)
    while True:
        if shell_export:
            # In shell-export mode, we read from /dev/tty to avoid messing with `eval`.
            with open('/dev/tty', 'r') as tty:
                print("> ", end='', file=sys.stderr)
                sys.stderr.flush()
                choice = tty.readline().strip().lower()
        else:
            choice = input("> ").strip().lower()

        if choice == 'q':
            return None
        try:
            idx = int(choice)
            if 1 <= idx <= len(matches):
                return matches[idx - 1]
        except ValueError:
            pass
        print("Invalid input. Please enter an integer from the list, or 'q' to quit.", file=sys.stderr)

def main():
    args = parse_args()
    verbosity = args.verbose

    # If no keys specified, default to "openai"
    if not args.keys:
        args.keys = "openai"

    # Parse the keys (comma separated)
    provider_list = [k.strip() for k in args.keys.split(',') if k.strip()]

    # Dictionary of chosen environment variables { canonical_var_name: value_from_user_choice }
    chosen_env = {}

    # For each requested key/provider
    for provider_str in provider_list:
        matches, prefix = find_envvars_for_provider(provider_str)
        selected_var = select_envvar_from_list(matches, prefix, verbosity,
                                               reveal_values=args.reveal_values,
                                               shell_export=args.shell_export)
        if selected_var:
            chosen_env[prefix] = os.environ[selected_var]

    # If the user wants shell export statements, print them and exit
    if args.shell_export:
        for k, v in chosen_env.items():
            print(f'export {k}="{v}"')
        sys.exit(0)

    # If no command is given (and no shell export requested), show usage and exit
    if not args.command:
        log_message("No command specified. Exiting.", 1, verbosity, "INFO")
        print("Usage: aikey -k openai,replicate [--] command [arguments]\n", file=sys.stderr)
        sys.exit(1)

    # Prepare environment for the subprocess
    new_env = dict(os.environ)
    new_env.update(chosen_env)

    # If we want to see debug info about final environment vars set
    log_message("Will run command with these newly set environment variables:", 2, verbosity, "INFO")
    for k, v in chosen_env.items():
        if args.reveal_values:
            log_message(f"{k}={v}", 2, verbosity, "INFO")
        else:
            log_message(f"{k}=[HIDDEN]", 2, verbosity, "INFO")

    # Execute the command
    cmd_str = " ".join(args.command)
    log_message(f"Executing command: {cmd_str}", 1, verbosity, "INFO")
    result = None
    try:
        result = subprocess.run(args.command, env=new_env)
    except FileNotFoundError:
        log_message(f"Command not found: {args.command[0]}", 1, verbosity, "INFO")
        log_message("Attempting fallback by executing via interactive bash...", 1, verbosity, "INFO")
        # Build an export-like string to ensure that the environment variables chosen overwrite
        # the ones that might be loaded in an interactive shell.
        export_cmd_parts = []
        for k, v in chosen_env.items():
            export_cmd_parts.append(f'export {k}="{v}"')
        export_cmd = "; ".join(export_cmd_parts)
        fallback_cmd = ["bash", "-i", "-c", f'{export_cmd}; {cmd_str}']
        try:
            result = subprocess.run(fallback_cmd, env=new_env)
        except Exception as e:
            log_message(f"Error in fallback execution: {e}", 1, verbosity, "INFO")
            sys.exit(1)
    except Exception as e:
        log_message(f"Error running command: {e}", 1, verbosity, "INFO")
        sys.exit(1)

    if result is not None:
        sys.exit(result.returncode)
    else:
        sys.exit(127)

if __name__ == "__main__":
    main()