dotfiles

# Z.AI API usage module.
# Requires ZAI_API_KEY env var.

# Main entry point - show usage summary
export def main [] { usage }

# Show usage summary
export def usage [] {
	let token = $env | get --optional ZAI_API_KEY | if ($in | is-empty) { null } else { $in }
	if $token == null { error make { msg: "Missing ZAI_API_KEY", label: { text: "Set in secrets.nuon" } } }

	let res = try { http get --headers { Authorization: $'Bearer ($token)', Content-Type: application/json } https://api.z.ai/api/monitor/usage/quota/limit } catch { |_| error make { msg: "API error", label: { text: "check ZAI_API_KEY" } } }

	let name = { 3: "5 Hours", 5: "Monthly", 6: "Weekly" }

	let limits = $res.data.limits | each {|r|
		let ts = $r | get --optional nextResetTime | default null
		{ kind: $r.type, period: ($name | get --optional ($r.unit | into string) | default $r.unit), pct: ($r | get --optional percentage | default 0), current: ($r | get --optional currentValue | default null), usage: ($r | get --optional usage | default null), resets: (if $ts != null { $ts / 1000 | math round | into int | into string | into datetime -f "%s" | format date "%Y-%m-%d" } else { null }) }
	}

	let details = $res.data.limits | each {|r| $r | get --optional usageDetails } | compact | flatten | each {|d| {model: $d.modelCode, usage: $d.usage}}

	show $res.data.level $limits $details
}

# Return raw limits table
export def "usage limits" [] {
	let token = $env | get --optional ZAI_API_KEY | if ($in | is-empty) { null } else { $in }
	if $token == null { error make { msg: "Missing ZAI_API_KEY", label: { text: "Set in secrets.nuon" } } }

	let res = try { http get --headers { Authorization: $'Bearer ($token)', Content-Type: application/json } https://api.z.ai/api/monitor/usage/quota/limit } catch { |_| error make { msg: "API error", label: { text: "check ZAI_API_KEY" } } }

	let name = { 3: "5 Hours", 5: "Monthly", 6: "Weekly" }

	$res.data.limits | each {|r|
		let ts = $r | get --optional nextResetTime | default null
		{ kind: $r.type, period: ($name | get --optional ($r.unit | into string) | default $r.unit), pct: ($r | get --optional percentage | default 0), current: ($r | get --optional currentValue | default null), usage: ($r | get --optional usage | default null), resets: (if $ts != null { $ts / 1000 | math round | into int | into string | into datetime -f "%s" | format date "%Y-%m-%d" } else { null }) }
	}
}

# Return raw details table
export def "usage details" [] {
	let token = $env | get --optional ZAI_API_KEY | if ($in | is-empty) { null } else { $in }
	if $token == null { error make { msg: "Missing ZAI_API_KEY", label: { text: "Set in secrets.nuon" } } }

	let res = try { http get --headers { Authorization: $'Bearer ($token)', Content-Type: application/json } https://api.z.ai/api/monitor/usage/quota/limit } catch { |_| error make { msg: "API error", label: { text: "check ZAI_API_KEY" } } }

	$res.data.limits | each {|r| $r | get --optional usageDetails } | compact | flatten | each {|d| {model: $d.modelCode, usage: $d.usage}}
}

def show [plan: string, limits: table, details: table] {
	let c = (ansi c)
	let m = (ansi m)
	let b = (ansi b)
	let d = (ansi d)
	let e = (ansi rst)

	print $"($c)($plan) Plan($e)"

	let tokens = $limits | where kind == TOKENS_LIMIT
	if $tokens != [] { print ""; print $"($c)Tokens($e) LLM quota"; $tokens | each {|r| show_limit $r $d null } }

	let calls = $limits | where kind == TIME_LIMIT
	if $calls != [] { print ""; print $"($m)Service Calls($e)"; $calls | each {|r| show_limit $r $d $r.usage } }

	if $details != [] { print ""; print $"($b)By Service($e)"; $details | each {|d| print $"  ($d.model)($e)  ($d.usage)" } }; null
}

def show_limit [row: record, dim: string, total: any] {
	let bar = bar $row.pct
	let used = if $total != null { $" ($row.current)/($total)" } else { "" }
	let resets = if $row.resets != null { $" ($dim)resets ($row.resets)($dim)" } else { "" }
	print $"  ($bar) ($dim)($row.period)($used)($resets)"
}

def bar [pct: int] {
	let n = ($pct / 10 | math round | into int | if $in > 10 { 10 } else { if $in < 0 { 0 } else { $in } })
	let g = (ansi g)
	let d = (ansi d)
	let y = (ansi y)
	let r = (ansi r)
	let e = (ansi rst)
	let color = if $pct < 50 { $g } else if $pct < 80 { $y } else { $r }
	let filled = 1..$n | each { $"($g)▓" } | str join
	let empty = 1..(10 - $n) | each { $"($d)░" } | str join
	$"($filled)($empty)($e) ($color)($pct)%($e)"
}