Something went wrong. Try again.
zig library for building requests to query inference providers usage limits
Something went wrong. Try again.
7.0 kB · 164 lines
Zig
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165//! CLI over the core: for debugging, and for parity tests against the TypeScript//! reference.const std = @import("std");const usage = @import("usage.zig");
const help_text = \\usage <command> \\ providers supported providers info \\ advance <provider> <credential.json> <responses.json>... advance feed state across rounds \\ [--now-ms N] \\ advance - the same, from one JSON envelope line on stdin: \\ {"provider","credential","responses":[round...],"nowMs"?} \\ claims <access-token> account id and email from an OAuth token \\;
fn writeTo(file: std.Io.File, io: std.Io, bytes: []const u8) !void { var buffer: [8192]u8 = undefined; var writer = file.writer(io, &buffer); try writer.interface.writeAll(bytes); try writer.flush();}
fn emit(init: std.process.Init, bytes: []const u8) !void { try writeTo(std.Io.File.stdout(), init.io, bytes); try writeTo(std.Io.File.stdout(), init.io, "\n");}
fn fail(init: std.process.Init, message: []const u8) !u8 { try writeTo(std.Io.File.stderr(), init.io, message); try writeTo(std.Io.File.stderr(), init.io, "\n\n"); try writeTo(std.Io.File.stderr(), init.io, help_text); return 1;}
fn nowMs() i64 { var ts: std.posix.timespec = undefined; if (std.posix.errno(std.posix.system.clock_gettime(.REALTIME, &ts)) != .SUCCESS) return 0; return @as(i64, ts.sec) * 1000 + @divTrunc(@as(i64, ts.nsec), 1_000_000);}
fn readFile(init: std.process.Init, path: []const u8, arena: std.mem.Allocator) ![]u8 { return std.Io.Dir.cwd().readFileAlloc(init.io, path, arena, std.Io.Limit.limited(8 * 1024 * 1024));}
/// Replays every round through one fresh state, printing one step per line: a host keeps the/// responses, and the process keeps nothing.fn advanceRounds( init: std.process.Init, arena: std.mem.Allocator, provider: []const u8, credential: []const u8, rounds: []const []const u8, now_ms: i64,) !u8 { var state = usage.State.init(arena); defer state.deinit();
for (rounds) |responses| { var w = usage.Writer.init(arena); const outcome = usage.advance(arena, &state, provider, credential, responses, now_ms) catch |e| { var err_w = usage.Writer.init(arena); usage.writeStep(&err_w, .{ .err = @errorName(e) }) catch return 1; try emit(init, err_w.items()); return 0; }; try usage.writeStep(&w, outcome); try emit(init, w.items()); } return 0;}
/// `advance -`: the whole call as one JSON envelope on stdin, so a credential never reaches argv,/// where any local process can read it. The envelope ends at the first newline, not only at EOF:/// a host that cannot half-close a pipe (an Erlang port) writes one line and reads the answer.fn advanceEnvelope(init: std.process.Init, arena: std.mem.Allocator) !u8 { const buffer = try arena.alloc(u8, 8 * 1024 * 1024); var reader = std.Io.File.stdin().readerStreaming(init.io, buffer); const input = (reader.interface.takeDelimiter('\n') catch |e| return fail(init, @errorName(e))) orelse return fail(init, "the envelope is empty");
const parsed = std.json.parseFromSliceLeaky(std.json.Value, arena, input, .{}) catch return fail(init, "the envelope is not JSON"); const root = usage.asObject(parsed) orelse return fail(init, "the envelope is not an object"); const provider = usage.asString(root.get("provider")) orelse return fail(init, "the envelope needs a provider"); const credential = root.get("credential") orelse return fail(init, "the envelope needs a credential"); const responses = usage.asArray(root.get("responses")) orelse return fail(init, "the envelope needs responses"); const now_ms = usage.toInt(i64, usage.asNumber(root.get("nowMs"))) orelse nowMs();
// An envelope asking for the first requests has no rounds yet; that is one empty round. var rounds = std.ArrayList([]const u8).empty; if (responses.items.len == 0) try rounds.append(arena, "[]"); for (responses.items) |round| { try rounds.append(arena, try std.json.Stringify.valueAlloc(arena, round, .{})); } const credential_text = try std.json.Stringify.valueAlloc(arena, credential, .{}); return advanceRounds(init, arena, provider, credential_text, rounds.items, now_ms);}
pub fn main(init: std.process.Init) !u8 { const arena = init.arena.allocator(); const argv = try std.process.Args.toSlice(init.minimal.args, arena); if (argv.len < 2) return fail(init, "missing command");
const command = argv[1];
if (std.mem.eql(u8, command, "providers")) { var w = usage.Writer.init(arena); try usage.writeProviders(&w); try emit(init, w.items()); return 0; }
if (std.mem.eql(u8, command, "advance")) { if (argv.len == 3 and std.mem.eql(u8, argv[2], "-")) return advanceEnvelope(init, arena); if (argv.len < 5) return fail(init, "advance needs <provider> <credential.json> <responses.json>... or -"); const provider = argv[2]; const cred_arg = argv[3];
var resp_args = std.ArrayList([]const u8).empty; var now_ms = nowMs(); var i: usize = 4; while (i < argv.len) : (i += 1) { if (std.mem.eql(u8, argv[i], "--now-ms") and i + 1 < argv.len) { now_ms = std.fmt.parseInt(i64, argv[i + 1], 10) catch now_ms; i += 1; } else { try resp_args.append(arena, argv[i]); } }
if (resp_args.items.len == 0) return fail(init, "advance needs at least one responses argument");
const cred_str = if (std.mem.startsWith(u8, cred_arg, "{")) cred_arg else readFile(init, cred_arg, arena) catch |e| return fail(init, @errorName(e));
var rounds = std.ArrayList([]const u8).empty; for (resp_args.items) |resp_arg| { const resp_str = if (std.mem.startsWith(u8, resp_arg, "[") or std.mem.startsWith(u8, resp_arg, "{")) resp_arg else readFile(init, resp_arg, arena) catch |e| return fail(init, @errorName(e)); try rounds.append(arena, resp_str); } return advanceRounds(init, arena, provider, cred_str, rounds.items, now_ms); }
if (std.mem.eql(u8, command, "claims")) { if (argv.len < 3) return fail(init, "claims needs an access token"); var w = usage.Writer.init(arena); const outcome = try usage.claimsFromAccessToken(arena, argv[2], &w); switch (outcome) { .ok => { try emit(init, w.items()); return 0; }, .err => |text| return fail(init, text), } }
return fail(init, "unknown command");}