vertexLib.py

import json
import time
import requests
import jwt
import asyncio

# ─────────────────────────────────────────
# 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, str(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"]
    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",
    }
    resp = requests.post(token_url, data=data)
    if resp.status_code == 200:
        return True, resp.json()["access_token"]
    return False, resp.text


# ─────────────────────────────────────────
# 3.  Gemini 모델 목록 (필요 시 10개까지만)
# ─────────────────────────────────────────
def get_gemini_models(key, max_return: int = 10):
    url = (
        "https://generativelanguage.googleapis.com/v1beta/models"
        f"?key={key}&pageSize=1000"
    )
    resp = requests.get(url)
    if resp.status_code != 200:
        return ""
    models = resp.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 names


# ─────────────────────────────────────────
# 4.  Gemini 키 간단 상태 체크
# ─────────────────────────────────────────
def check_key_gemini_availability(key):
    """
    반환값: (bool, str)
      • True,  'ok'      : 정상
      • False, 'exceed'  : 사용량/쿼터 초과
      • False, 'invalid' : 잘못된 키/권한 없음
      • False, 'error'   : 기타 오류
    """
    # 방법 1) 모델 목록 시도
    model_list = get_gemini_models(key, 1)
    if model_list != "":
        return True, "ok"

    # 방법 2) 더미 요청 시도
    url = (
        "https://generativelanguage.googleapis.com/v1beta/models/"
        "gemini-1.5-flash:generateContent"
        f"?key={key}"
    )
    payload = {
        "contents": [{"role": "user", "parts": [{"text": ""}]}],
        "generationConfig": {"maxOutputTokens": 0},
    }
    try:
        resp = requests.post(url, headers={"Content-Type": "application/json"}, json=payload)
    except Exception:
        return False, "error"

    if resp.status_code == 200:
        return True, "ok"

    err = resp.json().get("error", {})
    code = err.get("code", 0)
    status = err.get("status", "")
    if status == "INVALID_ARGUMENT":
        return True, "ok"
    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 = (
        "https://generativelanguage.googleapis.com/v1beta/models/"
        f"{model}:generateContent?key={key}"
    )
    resp = requests.post(url, headers={"Content-Type": "application/json"}, json=payload)
    if resp.status_code == 200:
        return True, resp.json()
    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 resp:
        if resp.status != 200:
            return json.loads(await resp.text())
        return await resp.json()


# ─────────────────────────────────────────
# 7.  다중 키 간단 검사 예시 함수
# ─────────────────────────────────────────
async def quick_check(keys: list[str]):
    """
    인자로 받은 키들을 Google Gemini 키라고 가정하고
    (key, availability, status) 3-필드 리스트 반환.
    필요하다면 여기서 다른 클라우드 체크 루틴을 호출하도록 확장.
    """
    results = []
    for k in keys:
        ok, st = check_key_gemini_availability(k.strip())
        results.append({"key": k.strip(), "availability": ok, "status": st})
    return results


# ─────────────────────────────────────────
# 8.  단독 실행 테스트
# ─────────────────────────────────────────
if __name__ == "__main__":
    # 예시 키 목록 (각자 보유 키로 교체)
    test_keys = [
        "AIzaSyExampleINVALIDKEY_for_demo1",
        "AIzaSyExampleINVALIDKEY_for_demo2",
    ]
    # 동기 실행 예시
    final = asyncio.run(quick_check(test_keys))
    print(json.dumps(final, indent=2, ensure_ascii=False))