Update vertexLib.py

vertexLib.py

@@ -2,186 +2,108 @@ import json
 import time
 import requests
 import jwt
-import asyncio
 
-# ─────────────────────────────────────────
-# 1.  GCP 서비스계정 액세스 토큰
-# ─────────────────────────────────────────
 def get_access_token(client_email, private_key):
     current_time = int(time.time())
-    expiration_time = current_time + 600  # 10 분
-
+    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,
+    '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")
+        signed_jwt = jwt.encode(
+            claims,
+            private_key,
+            algorithm='RS256'
+        )
     except Exception as e:
-        return False, str(e)
-
+        return False, e
+        
     response = requests.post(
-        "https://oauth2.googleapis.com/token",
+        'https://oauth2.googleapis.com/token',
         data={
-            "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
-            "assertion": signed_jwt,
-        },
+            'grant_type': 'urn:ietf:params:oauth:grant-type:jwt-bearer',
+            'assertion': signed_jwt
+        }
     )
 
     if response.status_code == 200:
-        return True, response.json()["access_token"]
-    return False, response.text
+        return True, response.json()['access_token']
+    else:
+        return False, response.text
 
-
-# ─────────────────────────────────────────
-# 2.  GCP refresh-token 액세스 토큰
-# ─────────────────────────────────────────
 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",
+        "grant_type": "refresh_token"
     }
-    resp = requests.post(token_url, data=data)
-    if resp.status_code == 200:
-        return True, resp.json()["access_token"]
-    return False, resp.text
 
+    response = requests.post(token_url, data=data)
 
-# ─────────────────────────────────────────
-# 3.  Gemini 모델 목록 (필요 시 10개까지만)
-# ─────────────────────────────────────────
-def get_gemini_models(key, max_return: int = 10):
-    url = (
-        "https://generativelanguage.googleapis.com/v1beta/models"
-        f"?key={key}&pageSize=1000"
-    )
-    resp = requests.get(url)
-    if resp.status_code != 200:
-        return ""
-    models = resp.json().get("models", [])
-    names = [m["name"].split("/")[1] for m in models]
-    if len(names) > max_return:
-        return names[:max_return] + [f"...(+{len(names)-max_return})"]
-    return names
-
+    if response.status_code == 200:
+        return True, response.json()["access_token"]
+    else:
+        return False, response.text
 
-# ─────────────────────────────────────────
-# 4.  Gemini 키 간단 상태 체크
-# ─────────────────────────────────────────
-def check_key_gemini_availability(key):
-    """
-    반환값: (bool, str)
-      • True,  'ok'      : 정상
-      • False, 'exceed'  : 사용량/쿼터 초과
-      • False, 'invalid' : 잘못된 키/권한 없음
-      • False, 'error'   : 기타 오류
-    """
-    # 방법 1) 모델 목록 시도
-    model_list = get_gemini_models(key, 1)
-    if model_list != "":
-        return True, "ok"
+def get_gemini_models(key):
+    url_getListModel = f"https://generativelanguage.googleapis.com/v1beta/models?key={key}&pageSize=1000"
+    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 ""
 
-    # 방법 2) 더미 요청 시도
-    url = (
-        "https://generativelanguage.googleapis.com/v1beta/models/"
-        "gemini-1.5-flash:generateContent"
-        f"?key={key}"
-    )
+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},
+        "contents": [{ "role": "user", "parts": [{ "text": "" }] }],
+        "generationConfig": {
+            "maxOutputTokens": 0,
+        }
     }
     try:
-        resp = requests.post(url, headers={"Content-Type": "application/json"}, json=payload)
-    except Exception:
-        return False, "error"
-
-    if resp.status_code == 200:
-        return True, "ok"
-
-    err = resp.json().get("error", {})
-    code = err.get("code", 0)
-    status = err.get("status", "")
-    if status == "INVALID_ARGUMENT":
-        return True, "ok"
-    if code == 429 or status == "RESOURCE_EXHAUSTED":
-        return False, "exceed"
-    if code in (401, 403) or status in ("PERMISSION_DENIED", "UNAUTHENTICATED"):
-        return False, "invalid"
-    return False, "error"
-
-
-# ─────────────────────────────────────────
-# 5.  Gemini 실제 요청 (필요 시)
-# ─────────────────────────────────────────
-def send_gemini_request(key, payload, model: str = "gemini-1.5-flash"):
-    url = (
-        "https://generativelanguage.googleapis.com/v1beta/models/"
-        f"{model}:generateContent?key={key}"
-    )
-    resp = requests.post(url, headers={"Content-Type": "application/json"}, json=payload)
-    if resp.status_code == 200:
-        return True, resp.json()
-    return False, resp.text
-
+        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
 
-# ─────────────────────────────────────────
-# 6.  Vertex AI (Anthropic) 비동기 요청
-# ─────────────────────────────────────────
-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/"
-        f"{project_id}/locations/{region}/publishers/anthropic/models/{model}:streamRawPredict"
-    )
+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",
+        'Authorization': f'Bearer {access_token}',
+        'Content-Type': 'application/json; charset=utf-8'
     }
-    async with session.post(url=VERTEX_URL, headers=headers, data=payload) as resp:
-        if resp.status != 200:
-            return json.loads(await resp.text())
-        return await resp.json()
-
-
-# ─────────────────────────────────────────
-# 7.  다중 키 간단 검사 예시 함수
-# ─────────────────────────────────────────
-async def quick_check(keys: list[str]):
-    """
-    인자로 받은 키들을 Google Gemini 키라고 가정하고
-    (key, availability, status) 3-필드 리스트 반환.
-    필요하다면 여기서 다른 클라우드 체크 루틴을 호출하도록 확장.
-    """
-    results = []
-    for k in keys:
-        ok, st = check_key_gemini_availability(k.strip())
-        results.append({"key": k.strip(), "availability": ok, "status": st})
-    return results
-
 
-# ─────────────────────────────────────────
-# 8.  단독 실행 테스트
-# ─────────────────────────────────────────
-if __name__ == "__main__":
-    # 예시 키 목록 (각자 보유 키로 교체)
-    test_keys = [
-        "AIzaSyExampleINVALIDKEY_for_demo1",
-        "AIzaSyExampleINVALIDKEY_for_demo2",
-    ]
-    # 동기 실행 예시
-    final = asyncio.run(quick_check(test_keys))
-    print(json.dumps(final, indent=2, ensure_ascii=False))
+    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()