From f37f07e3503ee7d138596423a6d8eeaa5d171d1f Mon Sep 17 00:00:00 2001 From: based Date: Wed, 28 Aug 2024 05:55:31 +1000 Subject: [PATCH] long overdue --- APIKey.py | 1 + README.md | 2 +- VertexAI.py | 82 +++++++++++++++++++++++++++++++++++++++-------------- main.py | 19 +++++++------ 4 files changed, 73 insertions(+), 31 deletions(-) diff --git a/APIKey.py b/APIKey.py index 9571833..bbc6edc 100644 --- a/APIKey.py +++ b/APIKey.py @@ -50,6 +50,7 @@ class APIKey: elif provider == Provider.VERTEXAI: self.project_id = "" + self.has_opus = False elif provider == Provider.MISTRAL: self.subbed = False diff --git a/README.md b/README.md index afd11a8..6cec51c 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ Currently supports and validates keys for the services below, and checks for the - Google MakerSuite (List of available models + billing status) - AWS - (Admin status, auto-fetch the region, logging status, username, bedrock status + enabled models) - Azure - (Auto-fetch all deployments, auto-fetch best deployment/model, filter status, dall-e deployment) -- Google Cloud Vertex AI - (Requires a key file since oauth tokens expire hourly. Good luck scraping for those.) +- Google Cloud Vertex AI - (Requires a key file since oauth tokens expire hourly.) - MistralAI - (Subscription status) - OpenRouter - (Estimated balance, usage in $, credit limit, RPM, has purchased any credits) - ElevenLabs - (Key tier, remaining characters in plan, detect uncapped char quota, pro voice cloning limit, invoice details on pay as you go plans) diff --git a/VertexAI.py b/VertexAI.py index c5e07c3..6cd55b0 100644 --- a/VertexAI.py +++ b/VertexAI.py @@ -1,18 +1,15 @@ +import asyncio import APIKey import json -import vertexai -from google.cloud import aiplatform from google.oauth2 import service_account -import google.api_core.exceptions -from vertexai.language_models import TextGenerationModel +from google.auth.transport.requests import Request + +location = 'us-east5' +model = 'claude-3-opus@20240229' -location = 'us-central1' # location doesn't matter unlike azure/aws - - -def check_vertexai(key: APIKey): +async def check_vertexai(key: APIKey, session): try: - credentials = service_account.Credentials.from_service_account_file(key.api_key) with open(key.api_key, 'r') as file: data = json.load(file) if data.get('type') != 'service_account': @@ -23,25 +20,68 @@ def check_vertexai(key: APIKey): return key.project_id = project_id - aiplatform.init(credentials=credentials, location=location, project=key.project_id) - test_model_response(key, credentials) + credentials = service_account.Credentials.from_service_account_file( + key.api_key, + scopes=['https://www.googleapis.com/auth/cloud-platform'] + ) + credentials.refresh(Request()) - except google.api_core.exceptions.InvalidArgument: - key.api_key = f'"{key.api_key}"' - return True # if we get to the stage where google yells at us for a bad parameter, 99% sure the key works. - except Exception as e: + key.has_opus = await test_model_response(key, credentials.token, session) + return True + + except Exception: return -def test_model_response(key: APIKey, credentials): - vertexai.init(project=key.project_id, location=location, credentials=credentials) - model = TextGenerationModel.from_pretrained("text-bison@002") - model.predict("bweh", **{"temperature": 0.1, "max_output_tokens": 0}) +async def test_model_response(key: APIKey, access_token, session): + url = f"https://{location}-aiplatform.googleapis.com/v1/projects/{key.project_id}/locations/{location}/publishers/anthropic/models/{model}:rawPredict" + + headers = { + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json; charset=utf-8" + } + + data = { + "anthropic_version": "vertex-2023-10-16", + "messages": [ + {'role': 'user', 'content': ''} + ], + "max_tokens": 0, + } + + for i in range(5): + try: + async with session.post(url, headers=headers, json=data) as response: + resp = await response.json() + if response.status == 429: + print(f"Rate limited on vertexai model response, retrying in 5 seconds (attempt {i} of 5)") + await asyncio.sleep(5) + continue + # returned on non enabled models afaik, couldn't test on the anthropic models though since they're all enabled on my service account + elif response.status == 400 and resp.get('error', {}).get('status') == 'FAILED_PRECONDITION': + return False + elif response.status == 400 and resp.get('error', {}).get('type') == 'invalid_request_error': + return True + else: + return False + except Exception as e: + print(f"Error testing model response: {str(e)}") + return False def pretty_print_vertexai_keys(keys): print('-' * 90) print(f'Validated {len(keys)} Google Vertex AI keys:') - for key in keys: + + keys_with_opus = [key for key in keys if key.has_opus] + keys_without_opus = [key for key in keys if not key.has_opus] + + print(f'\nValid keys with Opus enabled: {len(keys_with_opus)}') + for key in keys_with_opus: + print(f'{key.api_key} | {key.project_id} | has opus') + + print(f'\nValid keys without Opus enabled: {len(keys_without_opus)}') + for key in keys_without_opus: print(f'{key.api_key} | {key.project_id}') - print(f'\n--- Total Valid Google Vertex AI Keys: {len(keys)} ---\n') + + print(f'\n--- Total Valid Google Vertex AI Keys: {len(keys)} ({len(keys_with_opus)} with Opus enabled) ---\n') diff --git a/main.py b/main.py index 6f2969c..1cdcfb9 100644 --- a/main.py +++ b/main.py @@ -171,13 +171,14 @@ def validate_azure(key: APIKey): api_keys.add(key) -def validate_vertexai(key: APIKey): - IO.conditional_print(f"Checking Vertex AI keyfile: {key.api_key}", args.verbose) - if check_vertexai(key) is None: - IO.conditional_print(f"Invalid Vertex AI keyfile: {key.api_key}", args.verbose) - return - IO.conditional_print(f"Vertex AI keyfile '{key.api_key}' is valid", args.verbose) - api_keys.add(key) +async def validate_vertexai(key: APIKey, sem): + async with sem, aiohttp.ClientSession() as session: + IO.conditional_print(f"Checking VertexAI keyfile: {key.api_key}", args.verbose) + if await check_vertexai(key, session) is None: + IO.conditional_print(f"Invalid VertexAI keyfile: {key.api_key}", args.verbose) + return + IO.conditional_print(f"VertexAI keyfile '{key.api_key}' is valid", args.verbose) + api_keys.add(key) async def execute_with_retries(func, key, sem, retries): @@ -222,7 +223,7 @@ async def validate_keys(): if not os.path.isfile(key): continue key_obj = APIKey(Provider.VERTEXAI, key) - futures.append(executor.submit(validate_vertexai, key_obj)) + tasks.append(execute_with_retries(validate_vertexai, key_obj, concurrent_connections, 5)) elif "sk-ant-" in key[:7]: match = anthropic_regex.match(key) if "ant-api03" in key else anthropic_secondary_regex.match(key) if not match: @@ -290,7 +291,7 @@ def get_invalid_keys(valid_oai_keys, valid_anthropic_keys, valid_ai21_keys, vali valid_makersuite_keys_set = set([key.api_key for key in valid_makersuite_keys]) valid_aws_keys_set = set([key.api_key for key in valid_aws_keys]) valid_azure_keys_set = set([key.api_key for key in valid_azure_keys]) - valid_vertexai_keys_set = set([key.api_key for key in valid_vertexai_keys]) + valid_vertexai_keys_set = set([f'"{key.api_key}"' for key in valid_vertexai_keys]) valid_mistral_keys_set = set([key.api_key for key in valid_mistral_keys]) valid_openrouter_keys_set = set([key.api_key for key in valid_openrouter_keys]) valid_elevenlabs_set = set([key.api_key for key in valid_elevenlabs_keys])