import os import sys import time import json import hmac import hashlib import datetime import requests from urllib.parse import urlparse # ========================================== # 🛠️ Configuration # ========================================== # Volcengine API Config ACCESS_KEY = os.getenv( "VOLC_ACCESS_KEY", "AKLTNGYxZmUyNzNmZjk3NGMyZmEzNjVjZmY4NGM5ZjcxMjU" ) SECRET_KEY = os.getenv( "VOLC_SECRET_KEY", "Tm1VNE5XSmlORGs1WXpGbU5EVmtabUUyT1RZeVltTTBNalkxTWpjeE1tTQ==" ) HOST = "open.volcengineapi.com" REGION = "cn-beijing" BUFFER = 300000 # Suffixes to clean from Volcengine response SUFFIXES_TO_REMOVE = [ "_data_collaboration_resource_pack", "协作奖励计划资源包", "免费在线推理资源包", "免费在线推理包", "资源包", ] # Volcengine Name -> LiteLLM Model Name Map VOLC_MODEL_MAP = { "doubao-seed-1-6": ["doubao-seed-1-6-250615", "doubao-seed-1-6-251015"], "doubao-seed-1-8": "doubao-seed-1-8-251228", "doubao-seed-2-0-pro": "doubao-seed-2-0-pro-260215", "doubao-seed-2-0-lite": "doubao-seed-2-0-lite-260215", "doubao-seed-code": "doubao-seed-code-preview-251028", "doubao-seed-translation": "doubao-seed-translation-250915", "deepseek-v3": "deepseek-v3-250324", "deepseek-v3-1": "deepseek-v3-1-terminus", "deepseek-v3-2": "deepseek-v3-2-251201", "deepseek-r1": "deepseek-r1-250528", "kimi-k2": "kimi-k2-thinking-251104", "glm-4_7": "glm-4-7-251222", } # ========================================== # 🧮 Utility Functions # ========================================== def normalize_volc_name(raw_name): """Volcengine API response cleaning.""" if not raw_name: return "unknown" clean_name = raw_name for suffix in sorted(SUFFIXES_TO_REMOVE, key=len, reverse=True): if clean_name.endswith(suffix): clean_name = clean_name.replace(suffix, "") break clean_name = clean_name.replace("_", "-").replace(".", "-") clean_name = clean_name.lower().strip() return VOLC_MODEL_MAP.get(clean_name, clean_name) def sign(key, msg): return hmac.new(key, msg.encode("utf-8"), hashlib.sha256).digest() def get_signature_key(key, date_stamp, region_name, service_name): k_date = sign(key.encode("utf-8"), date_stamp) k_region = sign(k_date, region_name) k_service = sign(k_region, service_name) k_signing = sign(k_service, "request") return k_signing def get_signed_headers(action, version, body_dict, service="billing"): t = datetime.datetime.utcnow() amz_date = t.strftime("%Y%m%dT%H%M%SZ") date_stamp = t.strftime("%Y%m%d") request_body = json.dumps(body_dict) payload_hash = hashlib.sha256(request_body.encode("utf-8")).hexdigest() canonical_headers = f"content-type:application/json\nhost:{HOST}\nx-content-sha256:{payload_hash}\nx-date:{amz_date}\n" signed_headers = "content-type;host;x-content-sha256;x-date" canonical_request = f"POST\n/\nAction={action}&Version={version}\n{canonical_headers}\n{signed_headers}\n{payload_hash}" credential_scope = f"{date_stamp}/{REGION}/{service}/request" string_to_sign = f"HMAC-SHA256\n{amz_date}\n{credential_scope}\n{hashlib.sha256(canonical_request.encode('utf-8')).hexdigest()}" signing_key = get_signature_key(SECRET_KEY, date_stamp, REGION, service) signature = hmac.new( signing_key, string_to_sign.encode("utf-8"), hashlib.sha256 ).hexdigest() authorization_header = f"HMAC-SHA256 Credential={ACCESS_KEY}/{credential_scope}, SignedHeaders={signed_headers}, Signature={signature}" return { "Content-Type": "application/json", "Host": HOST, "X-Date": amz_date, "X-Content-Sha256": payload_hash, "Authorization": authorization_header, }, request_body # ========================================== # 🔄 Fetch & Process Logic # ========================================== def fetch_all_packages(): print(f"Fetching packages from {HOST}...") action = "ListResourcePackages" version = "2022-01-01" all_packages = [] current_token = "" while True: try: payload = { "ResourceType": "Package", "MaxResults": "20", "NextToken": current_token, } headers, body = get_signed_headers(action, version, payload) url = f"https://{HOST}/?Action={action}&Version={version}" response = requests.post(url, data=body, headers=headers, timeout=60) response.raise_for_status() data = response.json() items = data.get("Result", {}).get("List", []) print(f"Fetched {len(items)} items in this page.") if items: all_packages.extend(items) next_token = data.get("Result", {}).get("NextToken", "") if not next_token or next_token == current_token: break current_token = next_token time.sleep(0.2) except Exception as e: print(f"Error fetching packages: {e}") break return all_packages def process_packages_breakdown(packages): model_counts = {} print("\n--- Raw Package Processing ---") for pkg in packages: raw_code = pkg.get("ConfigurationCode", "") raw_name = pkg.get("ConfigurationName", "Unknown") target_raw = raw_code if raw_code else raw_name model_name = normalize_volc_name(target_raw) status = pkg.get("Status", "") if status == "Effective": available = pkg.get("AvailableAmount", 0) else: available = 0 unit = pkg.get("Unit", "token") print( f"Pkg: Code='{raw_code}', Name='{raw_name}' -> Norm='{model_name}', Avail={available} {unit} (Effective: {status == 'Effective'})" ) try: if unit.lower() == "token": model_counts[model_name] = model_counts.get(model_name, 0) + int( float(available) ) except ValueError: continue return model_counts # ========================================== # ▶️ Main Execution # ========================================== if __name__ == "__main__": print(f"Debug Script for Volcengine Sync") print(f"Access Key: {ACCESS_KEY[:4]}...{ACCESS_KEY[-4:]}") packages = fetch_all_packages() print(f"\nTotal Packages Fetched: {len(packages)}") model_balances = process_packages_breakdown(packages) print("\n--- Final Model Balances ---") for model, balance in model_balances.items(): remaining = balance - BUFFER print(f"Model: {model}") print(f" Balance : {balance:,}") print(f" Buffer : {BUFFER:,}") print(f" Remaining: {remaining:,}") # Simulate Redis Keys targets = [] if model in VOLC_MODEL_MAP: val = VOLC_MODEL_MAP[model] targets = val if isinstance(val, list) else [val] if not targets: targets = [model] else: targets = [model] print(f" Redis Keys that would be set:") for t in targets: print(f" - limit:dynamic:volcengine:remaining:{t} = {int(remaining)}")