Something went wrong. Try again.
zig library for building requests to query inference providers usage limits
Something went wrong. Try again.
7.7 kB · 186 lines
Zig
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187//! Z.ai usage feed implementation.//!//! Mirrors prime-agent:packages/ai/src/utils/usage/providers/zai.ts field for field.const std = @import("std");const usage = @import("../usage.zig");
pub const ZAI_USAGE_URL = "https://api.z.ai/api/monitor/usage/quota/limit";
pub const ParseZaiResult = struct { planType: ?[]const u8 = null, limits: []const usage.UsageLimit = &[_]usage.UsageLimit{},};
pub fn parseZaiUsage(alloc: std.mem.Allocator, payload: ?std.json.Value) !ParseZaiResult { const root = usage.asObject(payload); const data = if (root) |r| usage.asObject(r.get("data")) else null; if (data == null) return ParseZaiResult{};
var limits = std.ArrayList(usage.UsageLimit).empty; const raw_limits = usage.asArray(data.?.get("limits"));
if (raw_limits) |arr| { for (arr.items) |entry| { const item = usage.asObject(entry) orelse continue;
const type_str = usage.asString(item.get("type")) orelse ""; const unit = usage.asNumber(item.get("unit")); const raw_pct = usage.asNumber(item.get("percentage")); const cur = usage.asNumber(item.get("currentValue")); const total = usage.asNumber(item.get("usage")) orelse usage.asNumber(item.get("number"));
var used_percent = raw_pct; if (used_percent == null and cur != null and total != null and total.? > 0.0) { used_percent = (cur.? / total.?) * 100.0; }
const raw_reset = usage.asNumber(item.get("nextResetTime")); const resets_at: ?i64 = if (raw_reset) |rr| (if (rr > 100_000_000_000.0) usage.toInt(i64, rr) else usage.toInt(i64, rr * 1000.0)) else null;
var label: []const u8 = undefined; var window_label: ?[]const u8 = null; var detail: ?[]const u8 = null;
// One type can repeat with a different period, so the period is part of the key. const number = usage.asNumber(item.get("number")); const unit_int = usage.toInt(i64, unit); const number_int = usage.toInt(i64, number); const base_id = if (type_str.len > 0) try std.ascii.allocLowerString(alloc, type_str) else "limit"; const id = if (unit_int != null and number_int != null) try std.fmt.allocPrint(alloc, "{s}:{d}:{d}", .{ base_id, unit_int.?, number_int.? }) else base_id; // Unit 3 counts hours; the other units are not pinned down, so they carry no length. const window_seconds: ?i64 = if (unit_int == 3 and number_int != null) number_int.? * 3600 else null;
if (std.mem.eql(u8, type_str, "TIME_LIMIT") or (unit != null and unit.? == 5.0)) { label = "Requests"; if (cur != null and total != null) { if (usage.toInt(i64, cur.?) != null and usage.toInt(i64, total.?) != null) { detail = try std.fmt.allocPrint(alloc, "{d} / {d} reqs", .{ usage.toInt(i64, cur.?).?, usage.toInt(i64, total.?).?, }); } } } else if (std.mem.eql(u8, type_str, "TOKENS_LIMIT") or (unit != null and (unit.? == 3.0 or unit.? == 6.0))) { label = "Tokens"; window_label = "5h window"; } else { label = if (type_str.len > 0) try usage.titleCase(alloc, type_str) else "Limit"; }
const status: []const u8 = if (used_percent != null and used_percent.? >= 100.0) "exhausted" else usage.statusFor(used_percent);
try limits.append(alloc, .{ .id = id, .label = label, .usedPercent = used_percent, .windowLabel = window_label, .windowSeconds = window_seconds, .resetsAt = resets_at, .detail = detail, .status = status, }); } }
const level = usage.asString(data.?.get("level")); const plan_type: ?[]const u8 = if (level) |l| blk: { const trimmed = std.mem.trim(u8, l, " \t\r\n"); if (trimmed.len == 0) break :blk null; break :blk try usage.titleCase(alloc, trimmed); } else null;
return ParseZaiResult{ .planType = plan_type, .limits = limits.items, };}
pub fn advance( alloc: std.mem.Allocator, credential_json: []const u8, responses_json: []const u8, _now_ms: i64,) !usage.Outcome { _ = _now_ms; const cred_parsed = std.json.parseFromSlice(std.json.Value, alloc, credential_json, .{}) catch { return .{ .report = .{ .error_message = "invalid credential JSON" } }; }; const cred_root = usage.asObject(cred_parsed.value) orelse { return .{ .report = .{ .error_message = "invalid credential JSON" } }; }; const kind = usage.asString(cred_root.get("kind")) orelse ""; if (!std.mem.eql(u8, kind, "apiKey")) { return .{ .report = .{ .error_message = "zai needs an apiKey credential" } }; }
const api_key = usage.asString(cred_root.get("apiKey")) orelse ""; const cred_account_id = usage.asString(cred_root.get("accountId")); const cred_email = usage.asString(cred_root.get("email"));
const resp_parsed = std.json.parseFromSlice(std.json.Value, alloc, responses_json, .{}) catch { return .{ .report = .{ .error_message = "invalid responses JSON" } }; }; const resp_array = usage.asArray(resp_parsed.value) orelse { return .{ .report = .{ .error_message = "invalid responses JSON" } }; };
if (resp_array.items.len == 0) { var headers = std.ArrayList([2][]const u8).empty; // zai requires authorization without Bearer prefix try headers.append(alloc, .{ "Authorization", api_key }); try headers.append(alloc, .{ "Accept", "application/json" }); try headers.append(alloc, .{ "User-Agent", "OpenCode-Status-Plugin/1.0" });
const requests = try alloc.alloc(usage.UsageRequest, 1); requests[0] = .{ .method = "GET", .url = ZAI_USAGE_URL, .headers = headers.items, }; return .{ .requests = requests }; }
const first = switch (try usage.parseFirstResponse(alloc, resp_array, "usage body is not JSON")) { .ok => |f| f, .err => |msg| return .{ .report = .{ .error_message = msg } }, }; const body_parsed = first.parsed;
const root = usage.asObject(body_parsed.value); const success_val = if (root) |r| r.get("success") else null; const is_success = if (success_val) |sv| switch (sv) { .bool => |b| b, else => true, } else true; const code_num = if (root) |r| usage.asNumber(r.get("code")) else null;
if (!is_success or (code_num != null and code_num.? != 200.0)) { const msg = if (root) |r| usage.asString(r.get("msg")) else null; if (msg) |m| { return .{ .report = .{ .error_message = m } }; } const err_msg = try std.fmt.allocPrint(alloc, "Z.ai returned code {d}", .{usage.toInt(i64, code_num) orelse 0}); return .{ .report = .{ .error_message = err_msg } }; }
const parsed = try parseZaiUsage(alloc, body_parsed.value);
return .{ .report = .{ .planType = parsed.planType, .accountId = cred_account_id, .email = cred_email, .limits = parsed.limits, .error_message = null, }, };}