vertexLib.py

import json
import time
import requests
import jwt

# ─────────────────────────────────────────
# 1.  GCP 서비스계정 액세스 토큰
# ─────────────────────────────────────────
def get_access_token(client_email, private_key):
    current_time = int(time.time())
    expiration_time = current_time + 600  # 10 분

    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, e

    response = requests.post(
        "https://oauth2.googleapis.com/token",
        data={
            "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


# ─────────────────────────────────────────
# 2.  GCP refresh-token 액세스 토큰
# ─────────────────────────────────────────
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",
    }

    response = requests.post(token_url, data=data)

    if response.status_code == 200:
        return True, response.json()["access_token"]
    else:
        return False, response.text


# ─────────────────────────────────────────
# 3.  (선택) 모델 목록 조회 – 최대 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:
        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.  “더미” 요청으로 키 상태만 판별
# ─────────────────────────────────────────
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 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 == "":
        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 요청
# ─────────────────────────────────────────
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)

    if resp.status_code == 200:
        return True, resp.json()
    else:
        return False, resp.text


# ─────────────────────────────────────────
# 6.  Vertex AI (Anthropic) 요청 – 비동기
# ─────────────────────────────────────────
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/"
        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=VERTEX_URL, headers=headers, data=payload) as response:
        if response.status != 200:
            return json.loads(await response.text())
        return await response.json()