From 16536c9bde6576dee25ffbcf00f69fc0166e4b71 Mon Sep 17 00:00:00 2001 From: based Date: Fri, 10 Nov 2023 07:57:50 +1000 Subject: [PATCH] added ai21 support --- AI21.py | 37 +++++++++++++++++++++++++++++++++++++ APIKey.py | 4 ++++ README.md | 4 ++-- main.py | 33 +++++++++++++++++++++++++-------- 4 files changed, 68 insertions(+), 10 deletions(-) create mode 100644 AI21.py diff --git a/AI21.py b/AI21.py new file mode 100644 index 0000000..a9e49b1 --- /dev/null +++ b/AI21.py @@ -0,0 +1,37 @@ +import requests +import APIKey + + +def check_ai21(key: APIKey): + + url = "https://api.ai21.com/studio/v1/j2-light/complete" + + payload = { + "prompt": "a", + "maxTokens": 1, + } + headers = { + "accept": "application/json", + "content-type": "application/json", + "Authorization": f"Bearer {key.api_key}" + } + + response = requests.post(url, json=payload, headers=headers) + if response.status_code not in [200, 402]: + return + + if response.status_code == 402: # unsure if this error code also applies to empty keys + key.trial_elapsed = True + + return True + + +def pretty_print_ai21_keys(keys): + print('-' * 90) + dead = 0 + print(f'Validated {len(keys)} AI21 keys:') + for key in keys: + if key.trial_elapsed: + dead += 1 + print(f'{key.api_key}' + (' | trial elapsed' if key.trial_elapsed else "")) + print(f'\n--- Total Valid AI21 Keys: {len(keys)} ({dead} trial elapsed, {len(keys) - dead} still valid) ---\n') diff --git a/APIKey.py b/APIKey.py index 2df5ca2..0a36c3f 100644 --- a/APIKey.py +++ b/APIKey.py @@ -18,7 +18,11 @@ class APIKey: self.pozzed = False self.rate_limited = False + elif provider == Provider.AI21: + self.trial_elapsed = False + class Provider(Enum): OPENAI = 1, ANTHROPIC = 2 + AI21 = 3 diff --git a/README.md b/README.md index 4b619a0..b6e013a 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # keychecker a bulk key checker for various AI services -currently supports openai and shows various attributes (rpm, org etc) and anthropic with pozzed status check +currently supports openai, anthropic and ai21 keys and lists various attributes for each (org, rpm limit, pozzed status etc) -may add support for aws, palm or ai21 keys in the future +may add support for aws or palm keys in the future diff --git a/main.py b/main.py index becfb56..4afdbe1 100644 --- a/main.py +++ b/main.py @@ -2,6 +2,7 @@ from time import sleep from Anthropic import check_anthropic, pretty_print_anthropic_keys from Logger import Logger from OpenAI import get_oai_model, get_oai_key_attribs, get_oai_org, pretty_print_oai_keys +from AI21 import check_ai21, pretty_print_ai21_keys from APIKey import APIKey, Provider from concurrent.futures import ThreadPoolExecutor, as_completed import sys @@ -10,7 +11,7 @@ import re api_keys = set() -print("Enter API keys (OpenAI/Anthropic) one per line. Press Enter on a blank line to start validation") +print("Enter API keys (OpenAI/Anthropic/AI21) one per line. Press Enter on a blank line to start validation") inputted_keys = set() while True: @@ -48,8 +49,15 @@ def validate_anthropic(key: APIKey, retry_count): api_keys.add(key) +def validate_ai21(key: APIKey): + if check_ai21(key) is None: + return + api_keys.add(key) + + oai_regex = re.compile('(sk-[A-Za-z0-9]{20}T3BlbkFJ[A-Za-z0-9]{20})') anthropic_regex = re.compile(r'sk-ant-api03-[A-Za-z0-9\-_]{93}AA') +ai21_regex = re.compile('[A-Za-z0-9]{32}') executor = ThreadPoolExecutor(max_workers=100) @@ -63,23 +71,29 @@ def validate_keys(): continue key_obj = APIKey(Provider.ANTHROPIC, key) futures.append(executor.submit(validate_anthropic, key_obj, 20)) - else: + elif "sk-" in key: match = oai_regex.match(key) if not match: continue key_obj = APIKey(Provider.OPENAI, key) futures.append(executor.submit(validate_openai, key_obj)) - + else: + match = ai21_regex.match(key) + if not match: + continue + key_obj = APIKey(Provider.AI21, key) + futures.append(executor.submit(validate_ai21, key_obj)) for _ in as_completed(futures): pass futures.clear() -def get_invalid_keys(valid_oai_keys, valid_anthropic_keys): +def get_invalid_keys(valid_oai_keys, valid_anthropic_keys, valid_ai21_keys): valid_oai_keys_set = set([key.api_key for key in valid_oai_keys]) valid_anthropic_keys_set = set([key.api_key for key in valid_anthropic_keys]) - invalid_keys = inputted_keys - valid_oai_keys_set - valid_anthropic_keys_set + valid_ai21_keys_set = set([key.api_key for key in valid_ai21_keys]) + invalid_keys = inputted_keys - valid_oai_keys_set - valid_anthropic_keys_set - valid_ai21_keys_set if len(invalid_keys) < 1: return print('\nInvalid Keys:') @@ -91,25 +105,28 @@ def output_keys(): validate_keys() valid_oai_keys = [] valid_anthropic_keys = [] + valid_ai21_keys = [] for key in api_keys: if key.provider == Provider.OPENAI: valid_oai_keys.append(key) elif key.provider == Provider.ANTHROPIC: valid_anthropic_keys.append(key) - + elif key.provider == Provider.AI21: + valid_ai21_keys.append(key) output_filename = "key_snapshots.txt" sys.stdout = Logger(output_filename) print("#" * 90) print(f"Key snapshot from {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") print("#" * 90) print(f'\n--- Checked {len(inputted_keys)} keys | {len(inputted_keys) - len(api_keys)} were invalid ---') - get_invalid_keys(valid_oai_keys, valid_anthropic_keys) # just for completeness’s sake + get_invalid_keys(valid_oai_keys, valid_anthropic_keys, valid_ai21_keys) # just for completeness’s sake print() if valid_oai_keys: pretty_print_oai_keys(valid_oai_keys) if valid_anthropic_keys: pretty_print_anthropic_keys(valid_anthropic_keys) - + if valid_ai21_keys: + pretty_print_ai21_keys(valid_ai21_keys) sys.stdout.file.close()