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, } try: signed_jwt = jwt.encode(claims, private_key, algorithm="RS256") except Exception as e: 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, }, ) return _classify_response(resp), resp.json().get("access_token", resp.text) # ───────────────────────────────────────── # 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", }, ) return _classify_response(resp), resp.json().get("access_token", resp.text) # ───────────────────────────────────────── # 3) Gemini 모델 목록 (간소화) # 성공 여부만 반환 # ───────────────────────────────────────── 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) # ───────────────────────────────────────── # 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}, } try: 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", } async with session.post(url, headers=headers, data=payload) as resp: text = await resp.text() return _classify_async_status(resp.status, text), text