diff --git a/pyproject.toml b/pyproject.toml index ae2e0553..fb83518a 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,9 @@ vllm = [ all = ["nnsight[test,diffusers,vllm]"] +[project.scripts] +nnsight = "nnsight.login:main" + [project.urls] Homepage = "https://github.com/ndif-team/nnsight" Website = "https://nnsight.net/" diff --git a/src/nnsight/__init__.py b/src/nnsight/__init__.py index c6e3f3c6..6960fea6 100755 --- a/src/nnsight/__init__.py +++ b/src/nnsight/__init__.py @@ -56,6 +56,7 @@ from .intervention.tracing.globals import save from .ndif import * +from .login import login from IPython import get_ipython diff --git a/src/nnsight/login.py b/src/nnsight/login.py new file mode 100644 index 00000000..8f732927 --- /dev/null +++ b/src/nnsight/login.py @@ -0,0 +1,70 @@ +"""Login helpers for storing your NDIF API key. + +This module exposes two entry points: + +- ``nnsight.login()`` for use in a notebook or interactive session. It prompts + for the API key with ``getpass`` so the key is never echoed to the screen and + saves it through the existing config machinery. +- ``main()`` for the ``nnsight login`` console command, which does the same thing + from a terminal. + +Both persist the key by calling ``CONFIG.set_default_api_key``, which writes it +to ``config.yaml`` so future sessions pick it up automatically. +""" + +from getpass import getpass +from typing import Optional + + +def login(api_key: Optional[str] = None) -> None: + """Store your NDIF API key so future sessions can use it. + + Args: + api_key: The NDIF API key. If not given, you will be prompted for it. + The prompt uses ``getpass`` so the key is not echoed. + + Examples: + >>> from nnsight import login + >>> login() + Enter your NDIF API key: + NDIF API key saved. + """ + + from . import CONFIG + + if api_key is None: + api_key = getpass("Enter your NDIF API key: ") + + api_key = api_key.strip() + + if not api_key: + print("No API key provided. Nothing was saved.") + return + + CONFIG.set_default_api_key(api_key) + + print("NDIF API key saved.") + + +def main() -> None: + """Console entry point for the ``nnsight login`` command.""" + + import argparse + + parser = argparse.ArgumentParser( + prog="nnsight", + description="Command line tools for nnsight.", + ) + subparsers = parser.add_subparsers(dest="command") + subparsers.add_parser("login", help="Store your NDIF API key.") + + args = parser.parse_args() + + if args.command == "login": + login() + else: + parser.print_help() + + +if __name__ == "__main__": + main()