AgentsMeeting/gateway/scripts/usage_collector_kimi.py

#!/usr/bin/env python3
"""
usage_collector_kimi.py — Kimi 用量采集（CDP 浏览器自动化）
==========================================================
连接 Chrome CDP proxy (localhost:3456)，打开 Kimi 控制台，
提取用量数据，保存到 usage_stats_kimi.json。

调用方式: python usage_collector_kimi.py
"""
import json, os, sys, time, urllib.request, urllib.error
from pathlib import Path
from datetime import datetime, timezone

CDP = "http://localhost:3456"
KIMI_CONSOLE = "https://www.kimi.com/membership/subscription?tab=quota"

SCRIPT_DIR = Path(__file__).resolve().parent
GATEWAY_DIR = SCRIPT_DIR.parent
TEMP_DIR = GATEWAY_DIR / "temp"
os.makedirs(str(TEMP_DIR), exist_ok=True)

OUTPUT_FILE = TEMP_DIR / "usage_stats_kimi.json"
KEY_ID = "key7"
PROVIDER = "kimi"


def cdp_get(path, timeout=10):
    url = f"{CDP}{path}"
    try:
        with urllib.request.urlopen(url, timeout=timeout) as r:
            return json.loads(r.read().decode())
    except Exception as e:
        return {"error": str(e)}


def cdp_post(path, body="", timeout=15):
    url = f"{CDP}{path}"
    try:
        data = body.encode("utf-8") if isinstance(body, str) else body
        req = urllib.request.Request(url, data=data, method="POST")
        req.add_header("Content-Type", "text/plain")
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return r.read().decode()
    except Exception as e:
        return f"ERROR: {e}"


def cdp_eval(target_id, js, timeout=10):
    return cdp_post(f"/eval?target={target_id}", js, timeout)


def main():
    print("[kimi] Starting collection...")

    # 1. Find tab with the membership subscription page
    targets = cdp_get("/targets") or []
    kimi_tab = None
    for t in targets:
        url = t.get("url", "")
        if "membership/subscription" in url or "kimi.com/membership" in url:
            kimi_tab = t
            break

    if kimi_tab:
        target_id = kimi_tab.get("targetId") or kimi_tab.get("id")
        print(f"[kimi] Found quota tab: {target_id}")
        # Navigate if needed
        cur_url = kimi_tab.get("url", "")
        if "tab=quota" not in cur_url:
            print(f"[kimi] Navigating to quota tab...")
            cdp_get(f"/navigate?target={target_id}&url={KIMI_CONSOLE}")
            time.sleep(5)
    else:
        print(f"[kimi] No quota tab found, opening new tab...")
        result = cdp_get(f"/new?url={KIMI_CONSOLE}")
        if isinstance(result, dict):
            target_id = result.get("targetId") or result.get("id", "")
        else:
            target_id = str(result)
        print(f"[kimi] New tab: {target_id}")
        time.sleep(5)

    if not target_id:
        print("[kimi] ERROR: no target ID")
        return

    # 2. Wait for page load + extract data
    print("[kimi] Extracting usage data...")

    # Simple extraction: get all text, parse with Python
    js_get_text = "document.body ? document.body.innerText : ''"
    raw_text = cdp_eval(target_id, js_get_text)

    # CDP eval returns JSON: {"value":"text content"} — unwrap it
    try:
        parsed = json.loads(raw_text) if isinstance(raw_text, str) else raw_text
        text = parsed.get("value", "") if isinstance(parsed, dict) else str(raw_text)
    except (json.JSONDecodeError, AttributeError):
        text = str(raw_text)

    # Parse the text in Python
    import re

    # Membership page has three metrics:
    #   总用量 (monthly): N%, reset YYYY-MM-DD
    #   5小时用量 (5h): Code N%, reset MM-DD HH:MM
    #   7天用量 (7d): Code N%, reset MM-DD HH:MM
    # The percentages appear in order: total, 5h, 7d

    # Extract all percentage values with context
    pct_matches = list(re.finditer(r'(\d+(?:\.\d+)?)%', text))
    # Extract dates in format MM-DD HH:MM or YYYY-MM-DD
    date_matches = list(re.finditer(r'(\d{4}-\d{2}-\d{2})|(\d{2}-\d{2}\s+\d{2}:\d{2})', text))

    # Build list of (pct, date) pairs based on position in text
    entries = []
    for m in pct_matches:
        pct_val = float(m.group(1))
        pct_pos = m.start()
        # Find nearest date after this percentage
        best_date = None
        for d in date_matches:
            if d.start() > pct_pos:
                best_date = d.group(1) or d.group(2)
                break
        entries.append((pct_val, best_date, pct_pos))

    # Assign by order: first = total (monthly), second = 5h, third = 7d
    def parse_date(dstr):
        """Parse date string to seconds remaining."""
        if not dstr:
            return 0
        now = datetime.now(timezone.utc)
        try:
            if '-' in dstr and ':' not in dstr:
                # YYYY-MM-DD
                d = datetime.strptime(dstr, "%Y-%m-%d")
                return max(0, int((d.replace(tzinfo=timezone.utc) - now).total_seconds()))
            elif '-' in dstr and ':' in dstr:
                # MM-DD HH:MM — assume current year
                year = now.year
                d = datetime.strptime(dstr, "%m-%d %H:%M").replace(year=year, tzinfo=timezone.utc)
                if d < now:
                    d = d.replace(year=year + 1)
                return max(0, int((d - now).total_seconds()))
        except Exception:
            pass
        return 0

    def get_status(pct):
        if pct >= 95: return "rate-limited"
        if pct >= 80: return "warn"
        return "ok"

    # Build metrics
    total_pct = entries[0][0] if len(entries) > 0 else 0
    total_reset = parse_date(entries[0][1]) if len(entries) > 0 else 0
    five_hr_pct = entries[1][0] if len(entries) > 1 else 0
    five_hr_reset = parse_date(entries[1][1]) if len(entries) > 1 else 0
    seven_d_pct = entries[2][0] if len(entries) > 2 else 0
    seven_d_reset = parse_date(entries[2][1]) if len(entries) > 2 else 0

    print(f"[kimi] Monthly: {total_pct}% | 5h: {five_hr_pct}% | 7d: {seven_d_pct}%")

    # Also take a screenshot for debugging
    try:
        cdp_get(f"/screenshot?target={target_id}&file={TEMP_DIR}/kimi_console.png")
        print("[kimi] Screenshot saved")
    except:
        pass

    # Build output with correctly mapped metrics
    # rolling → 5小时用量, weekly → 7天用量, monthly → 总用量
    output = {
        "ok": True,
        "last_refresh_iso": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "provider": PROVIDER,
        "accounts": [{
            "key_id": KEY_ID,
            "label": "key7 (Kimi)",
            "provider": PROVIDER,
            "workspace_id": "kimi",
            "rolling": {
                "usage_percent": five_hr_pct,
                "reset_in_sec": five_hr_reset,
                "status": get_status(five_hr_pct),
            },
            "weekly": {
                "usage_percent": seven_d_pct,
                "reset_in_sec": seven_d_reset,
                "status": get_status(seven_d_pct),
            },
            "monthly": {
                "usage_percent": total_pct,
                "reset_in_sec": total_reset,
                "status": get_status(total_pct),
            },
            "session_expired": False,
            "last_update_iso": datetime.now(timezone.utc).isoformat(timespec="seconds"),
            "error": None,
        }],
    }

    with open(str(OUTPUT_FILE), "w", encoding="utf-8") as f:
        json.dump(output, f, ensure_ascii=False, indent=2)
    print(f"[kimi] Output saved to {OUTPUT_FILE}")


if __name__ == "__main__":
    main()