feat: multi-provider usage — Kimi collector via CDP, unified /api/usage, frontend grouping
This commit is contained in:
@@ -0,0 +1,173 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
usage_collector.py — 多 Provider 用量采集主控
|
||||
=============================================
|
||||
编排所有 provider 的采集，合并到统一的 usage_stats.json。
|
||||
|
||||
Provider:
|
||||
- opencode_go: 本地 HTTP/CDP 采集（usage_collector_ocg.py）
|
||||
- kimi: SSH 远程调用 Windows 采集器（usage_collector_kimi.py）
|
||||
|
||||
输出: gateway/temp/usage_stats.json（Dashboard /api/usage 读取）
|
||||
"""
|
||||
import json, os, sys, subprocess
|
||||
from pathlib import Path
|
||||
from datetime import datetime, timezone
|
||||
|
||||
SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
GATEWAY_DIR = SCRIPT_DIR.parent
|
||||
TEMP_DIR = GATEWAY_DIR / "temp"
|
||||
LOGS_DIR = GATEWAY_DIR / "logs"
|
||||
os.makedirs(str(TEMP_DIR), exist_ok=True)
|
||||
os.makedirs(str(LOGS_DIR), exist_ok=True)
|
||||
|
||||
OUTPUT_FILE = TEMP_DIR / "usage_stats.json"
|
||||
WINDOWS_HOST = "192.168.1.16"
|
||||
WINDOWS_PYTHON = r"C:\Users\hmo\AppData\Local\Programs\Python\Python310\python.exe"
|
||||
WINDOWS_SCRIPT = r"D:\F\NewI\opencode\daily-workspace\projects\AgentsMeeting\gateway\scripts\usage_collector_kimi.py"
|
||||
|
||||
|
||||
def log(msg):
|
||||
ts = datetime.now().strftime("%H:%M:%S")
|
||||
line = f"[{ts}] {msg}"
|
||||
print(line, flush=True)
|
||||
|
||||
|
||||
def collect_opencode_go():
|
||||
"""运行本地 OCG 采集器。"""
|
||||
log("OCG: starting...")
|
||||
ocg_script = SCRIPT_DIR / "usage_collector_ocg.py"
|
||||
if not ocg_script.exists():
|
||||
# Fallback: old name
|
||||
ocg_script = SCRIPT_DIR / "usage_collector.py"
|
||||
|
||||
try:
|
||||
r = subprocess.run(
|
||||
[sys.executable, str(ocg_script)],
|
||||
capture_output=True, text=True, timeout=120,
|
||||
cwd=str(SCRIPT_DIR),
|
||||
)
|
||||
if r.returncode != 0 and r.stderr:
|
||||
log(f"OCG: FAILED (rc={r.returncode})")
|
||||
log(f"OCG stderr: {r.stderr[-500:]}")
|
||||
else:
|
||||
log(f"OCG: done (rc={r.returncode})")
|
||||
except Exception as e:
|
||||
log(f"OCG: EXCEPTION: {e}")
|
||||
|
||||
|
||||
def collect_kimi():
|
||||
"""SSH 到 Windows 运行 Kimi 采集器，scp 结果回来。"""
|
||||
log("Kimi: SSH to Windows...")
|
||||
try:
|
||||
r = subprocess.run(
|
||||
["ssh", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=10",
|
||||
"-o", "BatchMode=yes", f"hmo@{WINDOWS_HOST}",
|
||||
f'"{WINDOWS_PYTHON}" "{WINDOWS_SCRIPT}"'],
|
||||
capture_output=True, text=True, timeout=45,
|
||||
)
|
||||
for line in r.stdout.splitlines()[-5:]:
|
||||
log(f"Kimi: {line.strip()}")
|
||||
if r.returncode != 0:
|
||||
log(f"Kimi: SSH FAILED (rc={r.returncode}): {r.stderr[-200:]}")
|
||||
# Try scp alternative
|
||||
r2 = subprocess.run(
|
||||
["scp", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=10",
|
||||
f"hmo@{WINDOWS_HOST}:{WINDOWS_SCRIPT}", "/dev/null"],
|
||||
capture_output=True, timeout=10,
|
||||
)
|
||||
if r2.returncode == 0:
|
||||
log("Kimi: SSH failed but scp works — check script path")
|
||||
except subprocess.TimeoutExpired:
|
||||
log("Kimi: SSH TIMEOUT")
|
||||
except Exception as e:
|
||||
log(f"Kimi: SSH EXCEPTION: {e}")
|
||||
|
||||
# Try to scp the output file from Windows
|
||||
kimi_output = TEMP_DIR / "usage_stats_kimi.json"
|
||||
try:
|
||||
r = subprocess.run(
|
||||
["scp", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=10",
|
||||
f"hmo@{WINDOWS_HOST}:" + WINDOWS_SCRIPT.replace("usage_collector_kimi.py", "") + "../temp/usage_stats_kimi.json",
|
||||
str(kimi_output)],
|
||||
capture_output=True, text=True, timeout=15,
|
||||
)
|
||||
if kimi_output.exists():
|
||||
log(f"Kimi: output scp'd ({kimi_output.stat().st_size} bytes)")
|
||||
else:
|
||||
log("Kimi: scp failed or file empty")
|
||||
except Exception as e:
|
||||
log(f"Kimi: scp EXCEPTION: {e}")
|
||||
|
||||
|
||||
def merge_all():
|
||||
"""合并所有 provider 的数据到 usage_stats.json。"""
|
||||
accounts = []
|
||||
|
||||
# 1. Read existing OCG cache (from usage_collector_ocg output)
|
||||
# The OCG collector writes directly to usage_stats.json or a cache
|
||||
cache_file = TEMP_DIR / "usage_stats.json"
|
||||
if cache_file.exists():
|
||||
try:
|
||||
with open(str(cache_file), "r", encoding="utf-8") as f:
|
||||
cached = json.load(f)
|
||||
for a in cached.get("accounts", []):
|
||||
a["provider"] = a.get("provider", "opencode_go")
|
||||
accounts.append(a)
|
||||
log(f"Merge: {len(accounts)} OCG accounts from cache")
|
||||
except Exception as e:
|
||||
log(f"Merge: failed to read cache: {e}")
|
||||
|
||||
# 2. Read Kimi output
|
||||
kimi_file = TEMP_DIR / "usage_stats_kimi.json"
|
||||
if kimi_file.exists():
|
||||
try:
|
||||
with open(str(kimi_file), "r", encoding="utf-8") as f:
|
||||
kimi_data = json.load(f)
|
||||
for a in kimi_data.get("accounts", []):
|
||||
a["provider"] = a.get("provider", "kimi")
|
||||
# Check if this key already exists (update instead of duplicate)
|
||||
existing = [i for i, x in enumerate(accounts) if x.get("key_id") == a.get("key_id")]
|
||||
if existing:
|
||||
accounts[existing[0]] = a
|
||||
else:
|
||||
accounts.append(a)
|
||||
log(f"Merge: +{len(kimi_data.get('accounts',[]))} Kimi accounts")
|
||||
except Exception as e:
|
||||
log(f"Merge: failed to read Kimi data: {e}")
|
||||
|
||||
# 3. Count by provider
|
||||
providers = {}
|
||||
for a in accounts:
|
||||
p = a.get("provider", "unknown")
|
||||
providers[p] = providers.get(p, 0) + 1
|
||||
|
||||
output = {
|
||||
"ok": True,
|
||||
"last_refresh_iso": datetime.now(timezone.utc).isoformat(timespec="seconds"),
|
||||
"account_count": len(accounts),
|
||||
"collected_count": sum(1 for a in accounts if a.get("error") is None),
|
||||
"providers": providers,
|
||||
"accounts": accounts,
|
||||
}
|
||||
|
||||
with open(str(OUTPUT_FILE), "w", encoding="utf-8") as f:
|
||||
json.dump(output, f, ensure_ascii=False, indent=2)
|
||||
log(f"Merge: saved {len(accounts)} accounts ({providers}) to {OUTPUT_FILE}")
|
||||
|
||||
|
||||
def main():
|
||||
log("=== Multi-provider collection started ===")
|
||||
|
||||
# Run collectors in sequence (OCG takes longest, Kimi needs CDP sync)
|
||||
collect_opencode_go()
|
||||
collect_kimi()
|
||||
|
||||
# Merge
|
||||
merge_all()
|
||||
|
||||
log("=== Collection complete ===")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user