Update vertexLib.py

vertexLib.py

@@ -3,7 +3,7 @@ import time
 import requests
 import jwt
 
-def get_access_token(project_id, client_email, private_key):
+def get_access_token(client_email, private_key):
     current_time = int(time.time())
     expiration_time = current_time + 600
     
@@ -24,7 +24,6 @@ def get_access_token(project_id, client_email, private_key):
     except Exception as e:
         return False, e
         
-    # Request the access token
     response = requests.post(
         'https://oauth2.googleapis.com/token',
         data={
@@ -34,11 +33,67 @@ def get_access_token(project_id, client_email, private_key):
     )
 
     if response.status_code == 200:
-        access_token = response.json()['access_token']
-        return True, access_token
+        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 = {
@@ -51,4 +106,4 @@ async def send_gcp_request(session, project_id, access_token, payload, region='u
             response_data = await response.text()
             return json.loads(response_data)
             
-        return await response.json() 
+        return await response.json()