Something went wrong. Try again.
zig library for building requests to query inference providers usage limits
Something went wrong. Try again.
7.8 kB · 239 lines
Swift
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240import CUsageimport Foundation
/// One quota window, matching the contract in packages/ai/docs/usage-feeds.md.public struct UsageLimit: Codable, Sendable, Equatable { /// Stable machine key within the provider; `label` is display text. public var id: String public var label: String public var usedPercent: Double? public var windowLabel: String? public var windowSeconds: Int64? public var resetsAt: Int64? /// The model family or feature this limit is restricted to; nil counts all usage. public var scope: String? public var status: String public var detail: String?}
/// What a provider's feed reports for one account.public struct UsageReport: Codable, Sendable, Equatable { public var provider: String? public var providerName: String? public var accountId: String? public var email: String? public var planType: String? public var limits: [UsageLimit] public var balance: String? public var notes: [String]? public var error: String?}
/// A provider with a feed, as the core describes it.public struct UsageProviderInfo: Codable, Sendable, Equatable { public var provider: String public var providerName: String public var credentialKinds: [String]}
/// The request the core wants sent. The host does the sending.////// Headers arrive as ordered pairs across the C ABI, because HTTP allows a name to repeat and a/// dictionary cannot. `headers` is the convenience view for the common case.public struct UsageRequest: Codable, Sendable, Equatable { public var kind: String = "http"
// HTTP request fields public var method: String = "GET" public var url: String = "" public var headerPairs: [[String]]? public var body: String?
// Command request fields public var command: String? public var args: [String]? public var stdin: String? public var timeoutMs: Int?
public var headers: [String: String] { Dictionary( (headerPairs ?? []).compactMap { pair in pair.count == 2 ? (pair[0], pair[1]) : nil }, uniquingKeysWith: { _, last in last } ) }
enum CodingKeys: String, CodingKey { case kind, method, url, body, command, args, stdin, timeoutMs case headerPairs = "headers" }
public init( kind: String = "http", method: String = "GET", url: String = "", headerPairs: [[String]]? = nil, body: String? = nil, command: String? = nil, args: [String]? = nil, stdin: String? = nil, timeoutMs: Int? = nil ) { self.kind = kind self.method = method self.url = url self.headerPairs = headerPairs self.body = body self.command = command self.args = args self.stdin = stdin self.timeoutMs = timeoutMs }
public init(from decoder: Decoder) throws { let container = try decoder.container(keyedBy: CodingKeys.self) self.kind = try container.decodeIfPresent(String.self, forKey: .kind) ?? "http" self.method = try container.decodeIfPresent(String.self, forKey: .method) ?? "GET" self.url = try container.decodeIfPresent(String.self, forKey: .url) ?? "" self.headerPairs = try container.decodeIfPresent([[String]].self, forKey: .headerPairs) self.body = try container.decodeIfPresent(String.self, forKey: .body) self.command = try container.decodeIfPresent(String.self, forKey: .command) self.args = try container.decodeIfPresent([String].self, forKey: .args) self.stdin = try container.decodeIfPresent(String.self, forKey: .stdin) self.timeoutMs = try container.decodeIfPresent(Int.self, forKey: .timeoutMs) }}
/// One response handed back to the core.public struct UsageResponse: Codable, Sendable, Equatable { public var status: Int public var body: String
public init(status: Int, body: String) { self.status = status self.body = body }}
/// Either the next requests to send, or the finished report.public enum UsageStep: Sendable { case requests([UsageRequest]) case report(UsageReport)
/// A provider that could not report says so in its report rather than by throwing. public var reportValue: UsageReport? { if case .report(let report) = self { return report } return nil }}
/// Opaque cross-round state managed by the host.public final class UsageState: @unchecked Sendable { let pointer: OpaquePointer?
public init() { self.pointer = usage_state_new() }
deinit { if let pointer { usage_state_free(pointer) } }}
public enum UsageCoreError: Error, CustomStringConvertible { case missingLibraryResult case invalidJSON(String) case providerRefused(String) case incompatibleAbiVersion(expected: UInt32, actual: UInt32)
public var description: String { switch self { case .missingLibraryResult: return "the usage core returned nothing" case .invalidJSON(let text): return "unexpected JSON from the usage core: \(text)" case .providerRefused(let message): return message case .incompatibleAbiVersion(let expected, let actual): return "incompatible usage core ABI version: expected \(expected), got \(actual)" } }}
public enum UsageCore { public static let expectedAbiVersion: UInt32 = 4 public static var abiVersion: UInt32 { usage_abi_version() }
public static func checkAbiVersion() throws { let version = abiVersion guard version == expectedAbiVersion else { throw UsageCoreError.incompatibleAbiVersion(expected: expectedAbiVersion, actual: version) } }
public static func providers() throws -> [UsageProviderInfo] { try checkAbiVersion() let json = try consume(usage_providers()) return try JSONDecoder().decode([UsageProviderInfo].self, from: Data(json.utf8)) }
/// One step of the host loop: hand the responses you already have (none on the first call) and /// get back either the requests to send or the report. public static func advance( state: UsageState? = nil, provider: String, credentialJSON: String, responses: [UsageResponse], nowMs: Int64 ) throws -> UsageStep { try checkAbiVersion() let responsesJSON = String(data: try JSONEncoder().encode(responses), encoding: .utf8) ?? "[]"
let raw = try consume( provider.withCString { providerPtr in credentialJSON.withCString { credentialPtr in responsesJSON.withCString { responsesPtr in usage_advance( state?.pointer, providerPtr, strlen(providerPtr), credentialPtr, strlen(credentialPtr), responsesPtr, strlen(responsesPtr), nowMs ) } } } )
guard let data = raw.data(using: .utf8), let object = try JSONSerialization.jsonObject(with: data) as? [String: Any] else { throw UsageCoreError.invalidJSON(raw) } if object["requests"] != nil { return .requests(try JSONDecoder().decode(UsageRequestsEnvelope.self, from: data).requests) } if let error = object["error"] as? String { throw UsageCoreError.providerRefused(error) } return .report(try JSONDecoder().decode(UsageReportEnvelope.self, from: data).report) }
/// Account id and email from an OAuth access token, without a request. public static func credentialsFromAccessToken(_ jwt: String) throws -> [String: String] { let json = try consume( jwt.withCString { pointer in usage_credentials_from_access_token(pointer, strlen(pointer)) } ) guard let data = json.data(using: .utf8) else { throw UsageCoreError.invalidJSON(json) } return (try? JSONDecoder().decode([String: String].self, from: data)) ?? [:] }
private struct UsageRequestsEnvelope: Codable { var requests: [UsageRequest] } private struct UsageReportEnvelope: Codable { var report: UsageReport }
/// Every call hands back a string the caller owns. NULL means allocation failed, which for a /// quota read is a failure like any other, so it becomes an error instead of a crash. private static func consume(_ pointer: UnsafeMutablePointer<CChar>?) throws -> String { guard let pointer else { throw UsageCoreError.missingLibraryResult } defer { usage_free(pointer) } return String(cString: pointer) }}