vertexLib.py

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