Something went wrong. Try again.
fork of Kludex/zigai, ported to zig 0.17
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330//! Provider-neutral request and run usage accounting.
const std = @import("std");
/// One provider-native integer counter that does not have a first-class field./// Names and the enclosing slice are borrowed from their owner.pub const Detail = struct { name: []const u8, value: u64,};
/// An exact USD amount represented as billionths of one dollar.pub const Cost = struct { nano_usd: u64,
pub fn fromUsd(value: f64) error{InvalidCost}!Cost { if (!std.math.isFinite(value) or value < 0) return error.InvalidCost; const scaled = value * nano_usd_per_usd; if (scaled > @as(f64, @floatFromInt(std.math.maxInt(u64)))) return error.InvalidCost; return .{ .nano_usd = @intFromFloat(@round(scaled)) }; }
pub fn usd(self: Cost) f64 { return @as(f64, @floatFromInt(self.nano_usd)) / nano_usd_per_usd; }
pub const nano_usd_per_usd: u64 = 1_000_000_000;};
pub const CostSource = enum { provider, price_table, mixed,};
/// Usage reported for one successful provider response. Input and output/// totals include their cached, audio, and reasoning subsets.pub const RequestUsage = struct { input_tokens: u64 = 0, cache_write_tokens: u64 = 0, cache_read_tokens: u64 = 0, output_tokens: u64 = 0, reasoning_tokens: u64 = 0, input_audio_tokens: u64 = 0, cache_audio_read_tokens: u64 = 0, output_audio_tokens: u64 = 0, details: []const Detail = &.{}, cost: ?Cost = null, cost_source: ?CostSource = null, /// Set only for price-table estimates. The slice is borrowed from the /// selected table, which must outlive the run. cost_table_version: ?[]const u8 = null, /// End-to-end provider latency measured by the agent. Direct model calls /// leave this null because the provider does not own the clock. duration_ms: ?u64 = null,
pub fn totalTokens(self: RequestUsage) u64 { return self.input_tokens +| self.output_tokens; }
pub fn detail(self: RequestUsage, name: []const u8) ?u64 { for (self.details) |item| if (std.mem.eql(u8, item.name, name)) return item.value; return null; }
pub fn hasValues(self: RequestUsage) bool { return self.totalTokens() != 0 or self.cache_write_tokens != 0 or self.cache_read_tokens != 0 or self.reasoning_tokens != 0 or self.input_audio_tokens != 0 or self.cache_audio_read_tokens != 0 or self.output_audio_tokens != 0 or self.details.len != 0 or self.cost != null or self.duration_ms != null; }};
/// Usage accumulated across one complete or paused agent run. `details` is an/// aggregate by name. Its storage follows the enclosing result arena.pub const RunUsage = struct { input_tokens: u64 = 0, cache_write_tokens: u64 = 0, cache_read_tokens: u64 = 0, output_tokens: u64 = 0, reasoning_tokens: u64 = 0, input_audio_tokens: u64 = 0, cache_audio_read_tokens: u64 = 0, output_audio_tokens: u64 = 0, details: []const Detail = &.{}, cost: ?Cost = null, cost_source: ?CostSource = null, cost_table_version: ?[]const u8 = null, requests: usize = 0, tool_calls: usize = 0, request_duration_ms: u64 = 0, run_duration_ms: u64 = 0,
pub fn totalTokens(self: RunUsage) u64 { return self.input_tokens +| self.output_tokens; }
pub fn detail(self: RunUsage, name: []const u8) ?u64 { for (self.details) |item| if (std.mem.eql(u8, item.name, name)) return item.value; return null; }
/// Records one provider attempt, including attempts that later fail and /// are retried. A null duration means the run has no monotonic clock. pub fn recordRequest(self: *RunUsage, duration_ms: ?u64) error{UsageOverflow}!void { const requests = std.math.add(usize, self.requests, 1) catch return error.UsageOverflow; const request_duration_ms = if (duration_ms) |duration| std.math.add(u64, self.request_duration_ms, duration) catch return error.UsageOverflow else self.request_duration_ms; self.requests = requests; self.request_duration_ms = request_duration_ms; }
/// Adds one completed provider response. Detail storage is allocated from /// `allocator`; callers normally pass the arena that owns the run result. pub fn addRequest( self: *RunUsage, allocator: std.mem.Allocator, request: RequestUsage, ) (std.mem.Allocator.Error || error{UsageOverflow})!void { var next = self.*; try addTokenFields(&next, request); try addCost(allocator, &next, request.cost, request.cost_source, request.cost_table_version); next.details = try mergeDetails(allocator, self.details, request.details); self.* = next; }
/// Adds another run, including its request, tool, latency, and cost totals. pub fn addRun( self: *RunUsage, allocator: std.mem.Allocator, other: RunUsage, ) (std.mem.Allocator.Error || error{UsageOverflow})!void { var next = self.*; try addTokenFields(&next, other); next.requests = std.math.add(usize, next.requests, other.requests) catch return error.UsageOverflow; next.tool_calls = std.math.add(usize, next.tool_calls, other.tool_calls) catch return error.UsageOverflow; next.request_duration_ms = std.math.add(u64, next.request_duration_ms, other.request_duration_ms) catch return error.UsageOverflow; next.run_duration_ms = std.math.add(u64, next.run_duration_ms, other.run_duration_ms) catch return error.UsageOverflow; try addCost(allocator, &next, other.cost, other.cost_source, other.cost_table_version); next.details = try mergeDetails(allocator, self.details, other.details); self.* = next; }
/// Deep-copies all borrowed run-usage storage into `allocator`. pub fn dupe(self: RunUsage, allocator: std.mem.Allocator) std.mem.Allocator.Error!RunUsage { var copy = self; const details = try allocator.alloc(Detail, self.details.len); for (self.details, details) |item, *target| target.* = .{ .name = try allocator.dupe(u8, item.name), .value = item.value, }; copy.details = details; copy.cost_table_version = if (self.cost_table_version) |version| try allocator.dupe(u8, version) else null; return copy; }};
/// Compatibility name for the original per-response usage type.pub const Usage = RequestUsage;
fn addTokenFields(target: anytype, source: anytype) error{UsageOverflow}!void { inline for (.{ "input_tokens", "cache_write_tokens", "cache_read_tokens", "output_tokens", "reasoning_tokens", "input_audio_tokens", "cache_audio_read_tokens", "output_audio_tokens", }) |name| { @field(target, name) = std.math.add(u64, @field(target, name), @field(source, name)) catch return error.UsageOverflow; }}
fn mergeDetails( allocator: std.mem.Allocator, left: []const Detail, right: []const Detail,) (std.mem.Allocator.Error || error{UsageOverflow})![]const Detail { if (right.len == 0) return left; var merged: std.ArrayList(Detail) = .empty; defer merged.deinit(allocator); try merged.appendSlice(allocator, left); for (right) |incoming| { var found = false; for (merged.items) |*existing| { if (!std.mem.eql(u8, existing.name, incoming.name)) continue; existing.value = std.math.add(u64, existing.value, incoming.value) catch return error.UsageOverflow; found = true; break; } if (!found) try merged.append(allocator, .{ .name = try allocator.dupe(u8, incoming.name), .value = incoming.value, }); } return merged.toOwnedSlice(allocator);}
fn addCost( allocator: std.mem.Allocator, target: *RunUsage, incoming: ?Cost, source: ?CostSource, table_version: ?[]const u8,) (std.mem.Allocator.Error || error{UsageOverflow})!void { const value = incoming orelse return; if (target.cost) |current| { target.cost = .{ .nano_usd = std.math.add(u64, current.nano_usd, value.nano_usd) catch return error.UsageOverflow }; const same_source = target.cost_source == source; const same_version = optionalStringsEqual(target.cost_table_version, table_version); if (!same_source or !same_version) { target.cost_source = .mixed; target.cost_table_version = null; } return; } target.cost = value; target.cost_source = source; target.cost_table_version = if (table_version) |version| try allocator.dupe(u8, version) else null;}
fn optionalStringsEqual(left: ?[]const u8, right: ?[]const u8) bool { if (left == null or right == null) return left == null and right == null; return std.mem.eql(u8, left.?, right.?);}
test "request usage exposes inclusive totals and details" { const value = RequestUsage{ .input_tokens = 10, .cache_read_tokens = 4, .output_tokens = 5, .reasoning_tokens = 2, .details = &.{.{ .name = "provider_counter", .value = 7 }}, }; try std.testing.expectEqual(@as(u64, 15), value.totalTokens()); try std.testing.expectEqual(@as(u64, 7), value.detail("provider_counter").?); try std.testing.expect(value.detail("missing") == null); try std.testing.expect(value.hasValues()); try std.testing.expect(!(RequestUsage{}).hasValues());}
test "run usage aggregates counters details costs and latency" { var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit(); var total = RunUsage{ .details = &.{.{ .name = "shared", .value = 2 }} }; try total.recordRequest(9); try total.addRequest(arena.allocator(), .{ .input_tokens = 8, .output_tokens = 3, .details = &.{ .{ .name = "shared", .value = 4 }, .{ .name = "native", .value = 1 }, }, .cost = .{ .nano_usd = 50 }, .cost_source = .price_table, .cost_table_version = "test-v1", }); try std.testing.expectEqual(@as(usize, 1), total.requests); try std.testing.expectEqual(@as(u64, 9), total.request_duration_ms); try std.testing.expectEqual(@as(u64, 6), total.detail("shared").?); try std.testing.expectEqual(@as(u64, 1), total.detail("native").?); try std.testing.expectEqual(@as(u64, 50), total.cost.?.nano_usd); try std.testing.expectEqual(CostSource.price_table, total.cost_source.?); try std.testing.expectEqualStrings("test-v1", total.cost_table_version.?); try std.testing.expect(total.detail("missing") == null);
try total.addRun(arena.allocator(), .{ .cost = .{ .nano_usd = 5 }, .cost_source = .price_table, .cost_table_version = "test-v1", }); try std.testing.expectEqual(CostSource.price_table, total.cost_source.?); try total.addRun(arena.allocator(), .{ .input_tokens = 2, .details = &.{.{ .name = "third", .value = 3 }}, .cost = .{ .nano_usd = 25 }, .cost_source = .provider, .requests = 2, .tool_calls = 1, .request_duration_ms = 4, .run_duration_ms = 5, }); try std.testing.expectEqual(CostSource.mixed, total.cost_source.?); try std.testing.expect(total.cost_table_version == null); try std.testing.expectEqual(@as(usize, 3), total.requests); try std.testing.expectEqual(@as(usize, 1), total.tool_calls); const copied = try total.dupe(arena.allocator()); try std.testing.expectEqualStrings("third", copied.details[copied.details.len - 1].name);}
test "cost conversion rejects invalid values" { try std.testing.expectEqual(@as(u64, 1_250_000_000), (try Cost.fromUsd(1.25)).nano_usd); try std.testing.expectApproxEqAbs(@as(f64, 1.25), (Cost{ .nano_usd = 1_250_000_000 }).usd(), 0.000000001); try std.testing.expectError(error.InvalidCost, Cost.fromUsd(-1)); try std.testing.expectError(error.InvalidCost, Cost.fromUsd(std.math.nan(f64)));}
test "usage aggregation reports overflow" { var total = RunUsage{ .input_tokens = std.math.maxInt(u64) }; try std.testing.expectError(error.UsageOverflow, total.addRequest(std.testing.allocator, .{ .input_tokens = 1 })); total = .{ .requests = std.math.maxInt(usize) }; try std.testing.expectError(error.UsageOverflow, total.recordRequest(null)); total = .{ .cost = .{ .nano_usd = std.math.maxInt(u64) } }; try std.testing.expectError(error.UsageOverflow, total.addRequest(std.testing.allocator, .{ .cost = .{ .nano_usd = 1 }, }));}
test "usage detail aggregation cleans up allocation failures" { const Check = struct { fn run(allocator: std.mem.Allocator) !void { var arena = std.heap.ArenaAllocator.init(allocator); defer arena.deinit(); var total = RunUsage{ .details = &.{.{ .name = "one", .value = 1 }} }; try total.addRequest(arena.allocator(), .{ .details = &.{.{ .name = "two", .value = 2 }} }); } }; try std.testing.checkAllAllocationFailures(@import("testing.zig").non_resizing_allocator, Check.run, .{});}