Update vertexLib.py

vertexLib.py

@@ -2,6 +2,7 @@ import json
 import time
 import requests
 import jwt
+import asyncio
 
 # ─────────────────────────────────────────
 # 1.  GCP 서비스계정 액세스 토큰
@@ -21,7 +22,7 @@ def get_access_token(client_email, private_key):
     try:
         signed_jwt = jwt.encode(claims, private_key, algorithm="RS256")
     except Exception as e:
-        return False, e
+        return False, str(e)
 
     response = requests.post(
         "https://oauth2.googleapis.com/token",
@@ -33,8 +34,7 @@ def get_access_token(client_email, private_key):
 
     if response.status_code == 200:
         return True, response.json()["access_token"]
-    else:
-        return False, response.text
+    return False, response.text
 
 
 # ─────────────────────────────────────────
@@ -48,112 +48,100 @@ def get_access_token_refresh(client_id, client_secret, refresh_token):
         "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
+    resp = requests.post(token_url, data=data)
+    if resp.status_code == 200:
+        return True, resp.json()["access_token"]
+    return False, resp.text
 
 
 # ─────────────────────────────────────────
-# 3.  (선택) 모델 목록 조회 – 최대 10개만 반환
+# 3.  Gemini 모델 목록 (필요 시 10개까지만)
 # ─────────────────────────────────────────
-def get_gemini_models(key, max_return: int = 1):
-    """
-    모델이 너무 많아 가독성이 떨어지므로,
-    기본적으로 10개까지만 반환하고 나머지는 개수로 요약.
-    """
-    url = f"https://generativelanguage.googleapis.com/v1beta/models?key={key}&pageSize=1000"
-    response = requests.get(url)
-
-    if response.status_code != 200:
+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 = response.json().get("models", [])
+    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 None
+    if len(names) > max_return:
+        return names[:max_return] + [f"...(+{len(names)-max_return})"]
+    return names
 
 
 # ─────────────────────────────────────────
-# 4.  “더미” 요청으로 키 상태만 판별
+# 4.  Gemini 키 간단 상태 체크
 # ─────────────────────────────────────────
-def send_fake_gemini_request(key, model: str = "gemini-1.5-flash"):
+def check_key_gemini_availability(key):
     """
-    프롬프트를 빈 문자열로 보내 간단히 키 유효성을 체크.
-    반환값: error dict | '' | None
+    반환값: (bool, str)
+      • True,  'ok'      : 정상
+      • False, 'exceed'  : 사용량/쿼터 초과
+      • False, 'invalid' : 잘못된 키/권한 없음
+      • False, 'error'   : 기타 오류
     """
-    url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={key}"
+    # 방법 1) 모델 목록 시도
+    model_list = get_gemini_models(key, 1)
+    if model_list != "":
+        return True, "ok"
+
+    # 방법 2) 더미 요청 시도
+    url = (
+        "https://generativelanguage.googleapis.com/v1beta/models/"
+        "gemini-1.5-flash:generateContent"
+        f"?key={key}"
+    )
     payload = {
         "contents": [{"role": "user", "parts": [{"text": ""}]}],
         "generationConfig": {"maxOutputTokens": 0},
     }
     try:
         resp = requests.post(url, headers={"Content-Type": "application/json"}, json=payload)
-        return resp.json().get("error", "")
     except Exception:
-        return None
-
-
-def check_key_gemini_availability(key):
-    """
-    ▶ 반환: (bool, str)
-        • (True,  'ok')       : 키 유효, 사용 가능
-        • (False, 'exceed')   : quota / rate-limit 초과
-        • (False, 'invalid')  : 잘못된 키 혹은 권한 없음
-        • (False, 'error')    : 기타 네트워크/알 수 없는 오류
-    """
-    err = send_fake_gemini_request(key)
-
-    # 네트워크 실패
-    if err is None:
         return False, "error"
 
-    # 에러 필드가 없으면 정상
-    if err == "":
+    if resp.status_code == 200:
         return True, "ok"
 
-    # 에러 객체 분석
+    err = resp.json().get("error", {})
     code = err.get("code", 0)
     status = err.get("status", "")
-
-    # 빈 프롬프트 때문에 INVALID_ARGUMENT가 뜨는 경우 → 키는 정상
     if status == "INVALID_ARGUMENT":
         return True, "ok"
-
-    # quota 초과
     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 요청
+# 5.  Gemini 실제 요청 (필요 시)
 # ─────────────────────────────────────────
 def send_gemini_request(key, payload, model: str = "gemini-1.5-flash"):
-    url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={key}"
+    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()
-    else:
-        return False, resp.text
+    return False, resp.text
 
 
 # ─────────────────────────────────────────
-# 6.  Vertex AI (Anthropic) 요청 – 비동기
+# 6.  Vertex AI (Anthropic) 비동기 요청
 # ─────────────────────────────────────────
 async def send_gcp_request(
-    session, project_id, access_token, payload, region="us-east5", model="claude-3-5-sonnet@20240620"
+    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/"
@@ -163,8 +151,37 @@ async def send_gcp_request(
         "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()
 
-    async with session.post(url=VERTEX_URL, headers=headers, data=payload) as response:
-        if response.status != 200:
-            return json.loads(await response.text())
-        return await response.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))