vertexLib.py

import json
import time
import requests
import jwt

def get_access_token(client_email, private_key):
    current_time = int(time.time())
    expiration_time = current_time + 600
    
    claims = {
    'iss': client_email,
    'scope': 'https://www.googleapis.com/auth/cloud-platform',
    'aud': 'https://oauth2.googleapis.com/token',
    'exp': expiration_time,
    'iat': current_time
    }
    
    try:
        signed_jwt = jwt.encode(
            claims,
            private_key,
            algorithm='RS256'
        )
    except Exception as e:
        return False, e
        
    response = requests.post(
        'https://oauth2.googleapis.com/token',
        data={
            'grant_type': 'urn:ietf:params:oauth:grant-type:jwt-bearer',
            'assertion': signed_jwt
        }
    )

    if response.status_code == 200:
        return True, response.json()['access_token']
    else:
        return False, response.text

def get_access_token_refresh(client_id, client_secret, refresh_token):
    token_url = "https://oauth2.googleapis.com/token"

    data = {
        "client_id": client_id,
        "client_secret": client_secret,
        "refresh_token": refresh_token,
        "grant_type": "refresh_token"
    }

    response = requests.post(token_url, data=data)

    if response.status_code == 200:
        return True, response.json()["access_token"]
    else:
        return False, response.text

def get_gemini_models(key):
    url_getListModel = f"https://generativelanguage.googleapis.com/v1beta/models?key={key}"
    response = requests.get(url_getListModel)
    
    if response.status_code == 200:
        models = response.json()
        model_list = []
        for model in models['models']:
            #model_list[model['name'].split('/')[1]] = model['displayName']
            model_name = f"{model['name'].split('/')[1]}" # ({model['displayName']})"
            model_list.append(model_name)
        return model_list
    else:
        return ""

def send_fake_gemini_request(key, model='gemini-1.5-flash'):
    url = f'https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={key}'
    payload = {
        "contents": [{ "role": "user", "parts": [{ "text": "" }] }],
        "generationConfig": {
            "maxOutputTokens": 0,
        }
    }
    try:
        response = requests.post(url=url, headers={'Content-Type': 'application/json'}, json=payload)        
        err_msg = response.json().get('error', '')
        return err_msg
    except:
        return None
        
def send_gemini_request(key, payload, model='gemini-1.5-flash'):
    url = f'https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={key}'
    
    response = requests.post(url=url, headers={'Content-Type': 'application/json'}, json=payload)
    
    if response.status_code == 200:
        return True, response.json()
    else:
        return False, response.text

async def send_gcp_request(session, project_id, access_token, payload, region='us-east5', model='claude-3-5-sonnet@20240620'):
    VERTEX_URL = f'https://{region}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{region}/publishers/anthropic/models/{model}:streamRawPredict'    
    headers = {
        'Authorization': f'Bearer {access_token}',
        'Content-Type': 'application/json; charset=utf-8'
    }

    async with session.post(url=VERTEX_URL, headers=headers, data=payload) as response:
        if response.status != 200:
            response_data = await response.text()
            return json.loads(response_data)
            
        return await response.json()