From 681bdc3134c93bc432862024ec07bb8f7bdadb4c Mon Sep 17 00:00:00 2001 From: based Date: Thu, 15 Feb 2024 04:55:28 +1000 Subject: [PATCH] fix for openai gateway timeout ratelimit and some aws changes --- APIKey.py | 1 + AWS.py | 39 ++++++----------- OpenAI.py | 125 ++++++++++++++++++++++++++++-------------------------- main.py | 7 +-- 4 files changed, 84 insertions(+), 88 deletions(-) diff --git a/APIKey.py b/APIKey.py index 9f9c695..6b36174 100644 --- a/APIKey.py +++ b/APIKey.py @@ -28,6 +28,7 @@ class APIKey: self.admin_priv = False self.bedrock_enabled = False self.region = "" + self.alt_regions = [] self.useless_reasons = [] self.logged = False diff --git a/AWS.py b/AWS.py index 30dfa5f..fb4bf66 100644 --- a/AWS.py +++ b/AWS.py @@ -4,35 +4,13 @@ import APIKey import botocore.exceptions +# us-gov-west-1 is also a supported bedrock region but needs an additional security token aws_regions = [ - "us-east-2", "us-east-1", - "us-west-1", "us-west-2", - "af-south-1", - "ap-east-1", - "ap-south-2", - "ap-southeast-3", - "ap-southeast-4", - "ap-south-1", - "ap-northeast-3", - "ap-northeast-2", "ap-southeast-1", - "ap-southeast-2", "ap-northeast-1", - "ca-central-1", "eu-central-1", - "eu-west-1", - "eu-west-2", - "eu-south-1", - "eu-west-3", - "eu-south-2", - "eu-north-1", - "eu-central-2", - "il-central-1", - "me-south-1", - "me-central-1", - "sa-east-1" ] @@ -49,8 +27,9 @@ def check_aws(key: APIKey): bedrock_runtime_client = session.client("bedrock-runtime") region = get_region(session) - if region is not None: - key.region = region + if region is not None and len(region) > 0: + key.region = region[0] + key.alt_regions = region[1:] # key.bedrock_enabled = True key.useless = False else: @@ -102,6 +81,11 @@ def check_aws(key: APIKey): key.useless = False continue + # Admin keys will never expose this policy even if they are quarantined. + if "AWSCompromisedKeyQuarantine" in policy["PolicyName"]: + key.useless = True + key.useless_reasons.append('Quarantined Key') + break if not key.useless: check_logging(session, key) elif key.useless and policies is not None: @@ -115,6 +99,7 @@ def check_aws(key: APIKey): def get_region(session): + regions = [] for region in aws_regions: try: bedrock_client = session.client("bedrock", region_name=region) @@ -122,9 +107,10 @@ def get_region(session): cloudies = ['anthropic.claude-v1', 'anthropic.claude-v2'] models = [model['modelId'] for model in response.get('modelSummaries', [])] if all(model_id in models for model_id in cloudies): - return region + regions.append(region) except botocore.exceptions.ClientError: return + return regions def test_invoke_perms(bedrock_runtime_client): @@ -177,6 +163,7 @@ def pretty_print_aws_keys(keys): for key in ready_to_go_keys: print(f'{key.api_key}' + (f' | {key.username}' if key.username != "" else "") + (' | admin key' if key.admin_priv else "") + (f' | {key.region}' if key.region != "" else "") + + (f' | alt regions - {key.alt_regions}' if key.alt_regions else "") + (' | LOGGED KEY' if key.logged is True else "")) if needs_setup_keys: diff --git a/OpenAI.py b/OpenAI.py index 3b51093..e1c4c96 100644 --- a/OpenAI.py +++ b/OpenAI.py @@ -1,61 +1,66 @@ import APIKey +import asyncio oai_api_url = "https://api.openai.com/v1" oai_t1_rpm_limits = {"gpt-3.5-turbo": 3500, "gpt-4": 500, "gpt-4-32k": 20} oai_tiers = {40000: 'Free', 60000: 'Tier1', 80000: 'Tier2', 160000: 'Tier3', 1000000: 'Tier4', 2000000: 'Tier5'} -async def get_oai_model(key: APIKey, session): - async with session.get(f'{oai_api_url}/models', headers={'Authorization': f'Bearer {key.api_key}'}) as response: - if response.status != 200: - return - else: - data = await response.json() - models = data["data"] - top_model = "gpt-3.5-turbo" - for model in models: - if model["id"] == "gpt-4-32k": - top_model = model["id"] - break - elif model["id"] == "gpt-4": - top_model = model["id"] - key.model = top_model - return True - - -async def get_oai_key_attribs(key: APIKey, session): - chat_object = {"model": f'{key.model}', "messages": [{"role": "user", "content": ""}], "max_tokens": 0} - async with session.post(f'{oai_api_url}/chat/completions', - headers={'Authorization': f'Bearer {key.api_key}', 'accept': 'application/json'}, - json=chat_object) as response: - if response.status in [400, 429]: - data = await response.json() - message = data["error"]["type"] - if message is None: +async def get_oai_model(key: APIKey, session, retries): + for _ in range(retries): + async with session.get(f'{oai_api_url}/models', headers={'Authorization': f'Bearer {key.api_key}'}) as response: + if response.status == 200: + data = await response.json() + models = data["data"] + top_model = "gpt-3.5-turbo" + for model in models: + if model["id"] == "gpt-4-32k": + top_model = model["id"] + break + elif model["id"] == "gpt-4": + top_model = model["id"] + key.model = top_model + return True + elif response.status != 502: return - match message: - case "access_terminated": + await asyncio.sleep(0.5) + + +async def get_oai_key_attribs(key: APIKey, session, retries): + chat_object = {"model": f'{key.model}', "messages": [{"role": "user", "content": ""}], "max_tokens": 0} + for _ in range(retries): + async with session.post(f'{oai_api_url}/chat/completions', + headers={'Authorization': f'Bearer {key.api_key}', 'accept': 'application/json'}, + json=chat_object) as response: + if response.status in [400, 429]: + data = await response.json() + message = data["error"]["type"] + if message is None: return - case "billing_not_active": - return - case "insufficient_quota": - key.has_quota = False - case "invalid_request_error": - key.has_quota = True - key.rpm = int(response.headers.get("x-ratelimit-limit-requests")) - if key.rpm < oai_t1_rpm_limits[key.model]: # oddly seen some gpt4 trial keys - key.trial = True - key.tier = await get_oai_key_tier(key, session) - else: - return - return True + match message: + case "access_terminated": + return + case "billing_not_active": + return + case "insufficient_quota": + key.has_quota = False + case "invalid_request_error": + key.has_quota = True + key.rpm = int(response.headers.get("x-ratelimit-limit-requests")) + if key.rpm < oai_t1_rpm_limits[key.model]: # oddly seen some gpt4 trial keys + key.trial = True + key.tier = await get_oai_key_tier(key, session, retries) + return True + elif response.status != 502: + return + await asyncio.sleep(0.5) # this will weed out fake t4/t5 keys reporting a 10k rpm limit, those keys would have requested to have their rpm increased -async def get_oai_key_tier(key: APIKey, session): +async def get_oai_key_tier(key: APIKey, session, retries): if key.trial: return 'Free' chat_object = {"model": f'gpt-3.5-turbo', "messages": [{"role": "user", "content": ""}], "max_tokens": 0} - for _ in range(3): # we'll give it 3 shots if oai decides to shid itself. + for _ in range(retries): async with session.post(f'{oai_api_url}/chat/completions', headers={'Authorization': f'Bearer {key.api_key}', 'accept': 'application/json'}, json=chat_object) as response: @@ -64,25 +69,27 @@ async def get_oai_key_tier(key: APIKey, session): return oai_tiers[int(response.headers.get("x-ratelimit-limit-tokens"))] except (KeyError, TypeError, ValueError): continue - else: + elif response.status != 502: return + await asyncio.sleep(0.5) return +async def get_oai_org(key: APIKey, session, retries): + for _ in range(retries): + async with session.get(f'{oai_api_url}/organizations', headers={'Authorization': f'Bearer {key.api_key}'}) as response: + if response.status == 200: + data = await response.json() + orgs = data["data"] -async def get_oai_org(key: APIKey, session): - async with session.get(f'{oai_api_url}/organizations', headers={'Authorization': f'Bearer {key.api_key}'}) as response: - if response.status != 200: - return - - data = await response.json() - orgs = data["data"] - - for org in orgs: - if not org["personal"]: - if org["is_default"]: - key.default_org = org["name"] - key.organizations.append(org["name"]) - return True + for org in orgs: + if not org["personal"]: + if org["is_default"]: + key.default_org = org["name"] + key.organizations.append(org["name"]) + return True + elif response.status != 502: + return + await asyncio.sleep(0.5) def check_manual_increase(key: APIKey): diff --git a/main.py b/main.py index a3df788..b20789c 100644 --- a/main.py +++ b/main.py @@ -49,14 +49,15 @@ else: async def validate_openai(key: APIKey, sem): + retries = 10 async with sem, aiohttp.ClientSession() as session: IO.conditional_print(f"Checking OpenAI key: {key.api_key}", args.verbose) - if await get_oai_model(key, session) is None: + if await get_oai_model(key, session, retries) is None: IO.conditional_print(f"Invalid OpenAI key: {key.api_key}", args.verbose) return - if await get_oai_key_attribs(key, session) is None: + if await get_oai_key_attribs(key, session, retries) is None: return - if await get_oai_org(key, session) is None: + if await get_oai_org(key, session, retries) is None: return IO.conditional_print(f"OpenAI key '{key.api_key}' is valid", args.verbose) api_keys.add(key)