Update vertexLib.py

vertexLib.py

@@ -2,108 +2,142 @@ 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) 서비스 계정 → 액세스 토큰
+# ─────────────────────────────────────────
 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
+        "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, e
-        
-    response = requests.post(
-        'https://oauth2.googleapis.com/token',
+        return "False", f"JWT encode error: {e}"
+
+    resp = requests.post(
+        "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']
-    else:
-        return False, response.text
+    return _classify_response(resp), resp.json().get("access_token", resp.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"
-    }
+# ─────────────────────────────────────────
+#  2) 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",
+        },
+    )
 
-    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 모델 목록 (간소화)
+#     성공 여부만 반환
+# ─────────────────────────────────────────
 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 ""
-
-def send_fake_gemini_request(key, model='gemini-1.5-flash'):
-    url = f'https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={key}'
+    url = f"https://generativelanguage.googleapis.com/v1beta/models?key={key}&pageSize=1"
+    resp = requests.get(url)
+    return _classify_response(resp)
+
+
+# ─────────────────────────────────────────
+#  4) Gemini 더미 호출 (쿼터 체크)
+# ─────────────────────────────────────────
+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:
-        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'    
+        resp = requests.post(url, headers={"Content-Type": "application/json"}, json=payload)
+        return _classify_response(resp)
+    except Exception:
+        return "False"
+
+
+# ─────────────────────────────────────────
+#  5) Gemini 실제 호출 (요약 상태 반환)
+# ─────────────────────────────────────────
+def send_gemini_request(key, payload, model="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
+
+
+# ─────────────────────────────────────────
+#  6) Vertex AI + Anthropic (streamRawPredict)
+#     상태만 간소화해 반환
+# ─────────────────────────────────────────
+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",
+):
+    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 response:
-        if response.status != 200:
-            response_data = await response.text()
-            return json.loads(response_data)
-            
-        return await response.json()
+    async with session.post(url, headers=headers, data=payload) as resp:
+        text = await resp.text()
+        return _classify_async_status(resp.status, text), text