Update vertexLib.py

vertexLib.py

@@ -2,46 +2,13 @@ import json
 import time
 import requests
 import jwt
-import aiohttp
-
 
 # ─────────────────────────────────────────
-#  공통: 응답을 세 가지로 축약해서 표시
-#        "True"    → 정상
-#        "exceed"  → 쿼터‧레이트리밋 초과
-#        "False"   → 그 외 실패
-# ─────────────────────────────────────────
-def _classify_response(resp: requests.Response) -> str:
-    if resp.status_code == 200:
-        return "True"
-
-    try:
-        err = resp.json().get("error", {})
-        code = err.get("code", 0)
-        msg = str(err.get("message", "")).lower()
-
-        if code == 429 or "quota" in msg or "rate" in msg or "exceed" in msg:
-            return "exceed"
-    except Exception:
-        pass
-
-    return "False"
-
-
-def _classify_async_status(code: int, text: str) -> str:
-    if code == 200:
-        return "True"
-    if code == 429 or any(kw in text.lower() for kw in ("quota", "rate", "exceed")):
-        return "exceed"
-    return "False"
-
-
-# ─────────────────────────────────────────
-#  1) 서비스 계정 → 액세스 토큰
+# 1.  GCP 서비스계정 액세스 토큰
 # ─────────────────────────────────────────
 def get_access_token(client_email, private_key):
     current_time = int(time.time())
-    expiration_time = current_time + 600
+    expiration_time = current_time + 600  # 10 분
 
     claims = {
         "iss": client_email,
@@ -54,9 +21,9 @@ def get_access_token(client_email, private_key):
     try:
         signed_jwt = jwt.encode(claims, private_key, algorithm="RS256")
     except Exception as e:
-        return "False", f"JWT encode error: {e}"
+        return False, e
 
-    resp = requests.post(
+    response = requests.post(
         "https://oauth2.googleapis.com/token",
         data={
             "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
@@ -64,80 +31,140 @@ def get_access_token(client_email, private_key):
         },
     )
 
-    return _classify_response(resp), resp.json().get("access_token", resp.text)
+    if response.status_code == 200:
+        return True, response.json()["access_token"]
+    else:
+        return False, response.text
 
 
 # ─────────────────────────────────────────
-#  2) refresh_token → 액세스 토큰
+# 2.  GCP refresh-token 액세스 토큰
 # ─────────────────────────────────────────
 def get_access_token_refresh(client_id, client_secret, refresh_token):
-    resp = requests.post(
-        "https://oauth2.googleapis.com/token",
-        data={
-            "client_id": client_id,
-            "client_secret": client_secret,
-            "refresh_token": refresh_token,
-            "grant_type": "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)
 
-    return _classify_response(resp), resp.json().get("access_token", resp.text)
+    if response.status_code == 200:
+        return True, response.json()["access_token"]
+    else:
+        return False, response.text
 
 
 # ─────────────────────────────────────────
-#  3) Gemini 모델 목록 (간소화)
-#     성공 여부만 반환
+# 3.  (선택) 모델 목록 조회 – 최대 10개만 반환
 # ─────────────────────────────────────────
-def get_gemini_models(key):
-    url = f"https://generativelanguage.googleapis.com/v1beta/models?key={key}&pageSize=1"
-    resp = requests.get(url)
-    return _classify_response(resp)
+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:
+        return ""
+
+    models = response.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
 
 
 # ─────────────────────────────────────────
-#  4) Gemini 더미 호출 (쿼터 체크)
+# 4.  “더���” 요청으로 키 상태만 판별
 # ─────────────────────────────────────────
-def send_fake_gemini_request(key, model="gemini-1.5-flash"):
+def send_fake_gemini_request(key, model: str = "gemini-1.5-flash"):
+    """
+    프롬프트를 빈 문자열로 보내 간단히 키 유효성을 체크.
+    반환값: error dict | '' | None
+    """
     url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={key}"
     payload = {
         "contents": [{"role": "user", "parts": [{"text": ""}]}],
         "generationConfig": {"maxOutputTokens": 0},
     }
-
     try:
         resp = requests.post(url, headers={"Content-Type": "application/json"}, json=payload)
-        return _classify_response(resp)
+        return resp.json().get("error", "")
     except Exception:
-        return "False"
+        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 == "":
+        return True, "ok"
+
+    # 에러 객체 분석
+    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="gemini-1.5-flash"):
+def send_gemini_request(key, payload, model: str = "gemini-1.5-flash"):
     url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={key}"
     resp = requests.post(url, headers={"Content-Type": "application/json"}, json=payload)
-    return _classify_response(resp), resp.json() if resp.status_code == 200 else resp.text
+
+    if resp.status_code == 200:
+        return True, resp.json()
+    else:
+        return False, resp.text
 
 
 # ─────────────────────────────────────────
-#  6) Vertex AI + Anthropic (streamRawPredict)
-#     상태만 간소화해 반환
+# 6.  Vertex AI (Anthropic) 요청 – 비동기
 # ─────────────────────────────────────────
 async def send_gcp_request(
-    session: aiohttp.ClientSession,
-    project_id: str,
-    access_token: str,
-    payload: str,
-    region: str = "us-east5",
-    model: str = "claude-3-5-sonnet@20240620",
+    session, project_id, access_token, payload, region="us-east5", model="claude-3-5-sonnet@20240620"
 ):
-    url = f"https://{region}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{region}/publishers/anthropic/models/{model}:streamRawPredict"
+    VERTEX_URL = (
+        f"https://{region}-aiplatform.googleapis.com/v1/projects/"
+        f"{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, headers=headers, data=payload) as resp:
-        text = await resp.text()
-        return _classify_async_status(resp.status, text), text
+    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()