added ai21 support

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')

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

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

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()