Skip to content

Commit c26dcce

Browse files
committed
refactor(dk): replace jq with python3 for JSON parsing
- Add check_python3() function to verify Python 3 availability - Replace jq JSON parsing with Python 3 in fetch_usage() - Remove --http1.1 flag from curl (allow HTTP/2 auto-negotiation) - Update error type checks (jq_error -> python_error/parse_error) - Add check_python3 calls before list, current, run, use commands This removes the jq dependency, using only Python 3 which is pre-installed on macOS/Linux. Performance is unchanged.
1 parent 905a23f commit c26dcce

1 file changed

Lines changed: 84 additions & 35 deletions

File tree

bin/dk

Lines changed: 84 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@ EOF
3939

4040
die() { echo "错误: $*" >&2; exit 1; }
4141

42+
check_python3() {
43+
command -v python3 >/dev/null 2>&1 || die "需要 Python 3,请安装: https://www.python.org/downloads/"
44+
}
45+
4246
ensure_store() {
4347
mkdir -p "$DKM_HOME"
4448
touch "$KEYS_ENC_FILE"
@@ -188,7 +192,7 @@ fetch_usage() {
188192
local key="$1"
189193
local resp http_status body
190194
# 将 HTTP 状态码与 body 分开,非 200 直接标记为不可用(避免继续使用无效 key)
191-
if ! resp=$(curl -sS --http1.1 --retry 0 \
195+
if ! resp=$(curl -sS --retry 0 \
192196
--connect-timeout "$DKM_CURL_CONNECT_TIMEOUT" --max-time "$DKM_CURL_MAX_TIME" \
193197
-w '\n%{http_code}' -X GET 'https://app.factory.ai/api/organization/members/chat-usage' \
194198
-H "Authorization: Bearer ${key}" \
@@ -206,36 +210,80 @@ fetch_usage() {
206210
return
207211
fi
208212

209-
if ! printf "%s" "$body" | jq -r '
210-
def n(v): if v==null then "" else (v|tostring) end;
211-
def tot(s): if s==null then null else (s.totalAllowance // s.basicAllowance // s.allowance // null) end;
212-
def used(s): if s==null then null else (s.orgTotalTokensUsed // s.used // s.tokensUsed // 0) end;
213-
def over(s): if s==null then null else (s.orgOverageUsed // 0) end;
214-
(.usage) as $u |
215-
($u | (.endDate // .expire_at // .expires_at)) as $exp_raw |
216-
([ $u.standard, $u.premium, $u.total, $u.main ]
217-
| map(select(. != null) | {total: tot(.), used: ((used(.) // 0) + (over(.) // 0))})
218-
| map(select(.total != null))
219-
| (.[0] // {})) as $s |
220-
($s.total) as $total |
221-
($s.used) as $used |
222-
(if ($total==null or $used==null) then null else ($total - $used) end) as $remain |
223-
($exp_raw | if .==null then null else (try ((.|tonumber)/1000) catch .) end) as $exp_ts |
224-
{
225-
BALANCE: $remain,
226-
BALANCE_NUM: $remain,
227-
TOTAL: $total,
228-
USED: $used,
229-
EXPIRES: (if $exp_ts==null then ($exp_raw|tostring) else ( $exp_ts | strftime("%Y-%m-%d") ) end),
230-
EXPIRES_FULL: (if $exp_ts==null then ($exp_raw|tostring) else ( $exp_ts | strftime("%Y-%m-%d %H:%M:%S %Z") ) end),
231-
EXPIRES_TS: $exp_ts,
232-
RAW: ""
233-
}
234-
| to_entries
235-
| map("\(.key)=\(n(.value))")
236-
| .[]
237-
'; then
238-
printf "BALANCE=0\nBALANCE_NUM=0\nTOTAL=0\nUSED=0\nEXPIRES=?\nEXPIRES_FULL=?\nEXPIRES_TS=\nRAW=jq_error\n"
213+
if ! printf "%s" "$body" | python3 -c '
214+
import sys, json
215+
from datetime import datetime, timezone
216+
217+
def n(v):
218+
return "" if v is None else str(v)
219+
220+
def tot(s):
221+
if s is None:
222+
return None
223+
for k in ("totalAllowance", "basicAllowance", "allowance"):
224+
if k in s and s[k] is not None:
225+
return s[k]
226+
return None
227+
228+
def used(s):
229+
if s is None:
230+
return 0
231+
for k in ("orgTotalTokensUsed", "used", "tokensUsed"):
232+
if k in s and s[k] is not None:
233+
return s[k]
234+
return 0
235+
236+
def over(s):
237+
if s is None:
238+
return 0
239+
return s.get("orgOverageUsed") or 0
240+
241+
try:
242+
data = json.load(sys.stdin)
243+
u = data.get("usage", {})
244+
exp_raw = u.get("endDate") or u.get("expire_at") or u.get("expires_at")
245+
section = None
246+
for sec in (u.get("standard"), u.get("premium"), u.get("total"), u.get("main")):
247+
if sec is not None:
248+
t = tot(sec)
249+
if t is not None:
250+
section = {"total": t, "used": used(sec) + over(sec)}
251+
break
252+
total = section["total"] if section else None
253+
used_val = section["used"] if section else None
254+
remain = (total - used_val) if (total is not None and used_val is not None) else None
255+
exp_ts = None
256+
if exp_raw is not None:
257+
try:
258+
exp_ts = float(exp_raw) / 1000
259+
except (ValueError, TypeError):
260+
pass
261+
if exp_ts is not None:
262+
dt = datetime.fromtimestamp(exp_ts, tz=timezone.utc)
263+
expires = dt.strftime("%Y-%m-%d")
264+
expires_full = dt.strftime("%Y-%m-%d %H:%M:%S %Z")
265+
else:
266+
expires = str(exp_raw) if exp_raw else ""
267+
expires_full = str(exp_raw) if exp_raw else ""
268+
print(f"BALANCE={n(remain)}")
269+
print(f"BALANCE_NUM={n(remain)}")
270+
print(f"TOTAL={n(total)}")
271+
print(f"USED={n(used_val)}")
272+
print(f"EXPIRES={expires}")
273+
print(f"EXPIRES_FULL={expires_full}")
274+
print(f"EXPIRES_TS={n(exp_ts)}")
275+
print("RAW=")
276+
except Exception:
277+
print("BALANCE=0")
278+
print("BALANCE_NUM=0")
279+
print("TOTAL=0")
280+
print("USED=0")
281+
print("EXPIRES=?")
282+
print("EXPIRES_FULL=?")
283+
print("EXPIRES_TS=")
284+
print("RAW=parse_error")
285+
'; then
286+
printf "BALANCE=0\nBALANCE_NUM=0\nTOTAL=0\nUSED=0\nEXPIRES=?\nEXPIRES_FULL=?\nEXPIRES_TS=\nRAW=python_error\n"
239287
fi
240288
}
241289

@@ -530,7 +578,7 @@ auto_rotate_if_needed() {
530578
esac
531579
done <<<"$info"
532580

533-
if [[ "$raw" == http_* || "$raw" == curl_error || "$raw" == jq_error ]]; then
581+
if [[ "$raw" == http_* || "$raw" == curl_error || "$raw" == python_error || "$raw" == parse_error ]]; then
534582
return 1
535583
fi
536584
remain=$(calc_remain "$balnum" "$total" "$used")
@@ -876,6 +924,7 @@ cmd_use() {
876924
if [ $# -eq 0 ]; then
877925
local n=${#KEYS[@]}
878926
[ $n -gt 0 ] || die "暂无key"
927+
check_python3
879928
if ! cache_load_infos; then
880929
fetch_all_infos
881930
cache_save_infos
@@ -1198,10 +1247,10 @@ main() {
11981247
local cmd="${1:-help}"; shift || true
11991248
case "$cmd" in
12001249
add) cmd_add "$@";;
1201-
list|ls) cmd_list "$@";;
1202-
current) cmd_current "$@";;
1250+
list|ls) check_python3; cmd_list "$@";;
1251+
current) check_python3; cmd_current "$@";;
12031252
use) cmd_use "$@";;
1204-
run) cmd_run "$@";;
1253+
run) check_python3; cmd_run "$@";;
12051254
serve) cmd_serve "$@";;
12061255
rm|remove|del) cmd_rm "$@";;
12071256
uninstall) cmd_uninstall "$@";;

0 commit comments

Comments
 (0)