#!/usr/bin/env python3
"""Расход ключа шлюза AILab SCRUME (gw.ailabscrume.ru) — компактный JSON в stdout.
Вызывается плагином ailab-limit через shell.exec. Ключ берётся из .env профилей
и в рендерер не попадает. Один запрос: GET /key/info с пользовательским ключом
(master не нужен) — шлюз отдаёт spend, max_budget и budget_duration в рублях.
Формат ответа:
{"ok":true,"spend":5.21,"budget":100.0,"duration":"1d","pct":5.2,"alias":"testoff"}
{"ok":false,"err":"..."}
"""
import json
import os
import sys
import urllib.error
import urllib.request
BASE = "https://gw.ailabscrume.ru"
KEY_NAMES = (
"HERMES_CUSTOM_AILAB_API_KEY",
"CUSTOM_PROVIDER_AILAB_KEY",
"HERMES_CUSTOM_POLZA_AI_API_KEY",
"CUSTOM_PROVIDER_POLZA_AI_KEY",
"AILAB_API_KEY",
)
def candidates():
"""Пути к .env: корневой и все профили, активный — первым."""
root = os.path.expandvars(r"%LOCALAPPDATA%\hermes")
out = []
active = os.path.expandvars(r"%APPDATA%\hermes\active-profile.json")
if os.path.exists(active):
try:
with open(active, encoding="utf-8") as f:
name = json.load(f).get("profile") or json.load(f).get("active")
if name:
out.append(os.path.join(root, "profiles", name, ".env"))
except Exception:
pass
out.append(os.path.join(root, ".env"))
prof_dir = os.path.join(root, "profiles")
if os.path.isdir(prof_dir):
for name in sorted(os.listdir(prof_dir)):
out.append(os.path.join(prof_dir, name, ".env"))
seen, uniq = set(), []
for p in out:
if p not in seen:
seen.add(p)
uniq.append(p)
return uniq
def read_key(path):
"""Первая непустая переменная из KEY_NAMES в файле .env."""
try:
with open(path, encoding="utf-8", errors="ignore") as f:
lines = f.readlines()
except OSError:
return None, None
found = {}
for line in lines:
s = line.strip()
if not s or s.startswith("#") or "=" not in s:
continue
name, _, value = s.partition("=")
name = name.strip()
if name in KEY_NAMES:
value = value.strip().strip('"').strip("'")
if value.startswith("sk-"):
found[name] = value
for name in KEY_NAMES:
if name in found:
return name, found[name]
return None, None
def fetch(key, timeout=10.0):
req = urllib.request.Request(BASE + "/key/info", headers={"Authorization": "Bearer " + key})
with urllib.request.urlopen(req, timeout=timeout) as r:
data = json.loads(r.read())
return data.get("info", data)
def main():
for path in candidates():
name, key = read_key(path)
if not key:
continue
try:
info = fetch(key)
except urllib.error.HTTPError as e:
print(json.dumps({"ok": False, "err": "шлюз вернул %s" % e.code}, ensure_ascii=False))
return
except Exception as e:
print(json.dumps({"ok": False, "err": str(e)[:120]}, ensure_ascii=False))
return
spend = float(info.get("spend") or 0)
budget = info.get("max_budget")
budget = float(budget) if budget is not None else None
pct = round(spend / budget * 100, 1) if budget else None
print(json.dumps({
"ok": True,
"spend": round(spend, 4),
"budget": round(budget, 2) if budget else None,
"duration": info.get("budget_duration") or "",
"pct": pct,
"alias": info.get("key_alias") or "",
}, ensure_ascii=False))
return
print(json.dumps({"ok": False, "err": "ключ шлюза не найден в .env"}, ensure_ascii=False))
if __name__ == "__main__":
sys.exit(main())