mirror of
https://github.com/cunnymessiah/keychecker.git
synced 2026-05-11 02:40:13 -07:00
based
48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
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
|
|
|
|
|
|
location = 'us-central1' # location doesn't matter unlike azure/aws
|
|
|
|
|
|
def check_vertexai(key: APIKey):
|
|
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':
|
|
return
|
|
|
|
project_id = data.get('project_id')
|
|
if not project_id:
|
|
return
|
|
key.project_id = project_id
|
|
|
|
aiplatform.init(credentials=credentials, location=location, project=key.project_id)
|
|
test_model_response(key, credentials)
|
|
|
|
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:
|
|
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})
|
|
|
|
|
|
def pretty_print_vertexai_keys(keys):
|
|
print('-' * 90)
|
|
print(f'Validated {len(keys)} Google Vertex AI keys:')
|
|
for key in keys:
|
|
print(f'{key.api_key} | {key.project_id}')
|
|
print(f'\n--- Total Valid Google Vertex AI Keys: {len(keys)} ---\n')
|