174 lines
6.1 KiB
Python
174 lines
6.1 KiB
Python
#!/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()
|