293 lines
10 KiB
Swift
293 lines
10 KiB
Swift
import Foundation
|
|
#if canImport(FoundationNetworking)
|
|
import FoundationNetworking
|
|
#endif
|
|
|
|
public enum VertexAIFetchError: LocalizedError, Sendable {
|
|
case unauthorized
|
|
case forbidden
|
|
case noProject
|
|
case networkError(Error)
|
|
case invalidResponse(String)
|
|
case noData
|
|
|
|
public var errorDescription: String? {
|
|
switch self {
|
|
case .unauthorized:
|
|
"Vertex AI request unauthorized. Run `gcloud auth application-default login`."
|
|
case .forbidden:
|
|
"Access forbidden. Check your IAM permissions for Cloud Monitoring."
|
|
case .noProject:
|
|
"No Google Cloud project configured. Run `gcloud config set project PROJECT_ID`."
|
|
case let .networkError(error):
|
|
"Vertex AI network error: \(error.localizedDescription)"
|
|
case let .invalidResponse(message):
|
|
"Vertex AI response was invalid: \(message)"
|
|
case .noData:
|
|
"No Vertex AI usage data found for the current project."
|
|
}
|
|
}
|
|
}
|
|
|
|
public struct VertexAIUsageResponse: Sendable {
|
|
public let requestsUsedPercent: Double
|
|
public let tokensUsedPercent: Double?
|
|
public let resetsAt: Date?
|
|
public let resetDescription: String?
|
|
public let rawData: String?
|
|
|
|
public init(
|
|
requestsUsedPercent: Double,
|
|
tokensUsedPercent: Double?,
|
|
resetsAt: Date?,
|
|
resetDescription: String?,
|
|
rawData: String?)
|
|
{
|
|
self.requestsUsedPercent = requestsUsedPercent
|
|
self.tokensUsedPercent = tokensUsedPercent
|
|
self.resetsAt = resetsAt
|
|
self.resetDescription = resetDescription
|
|
self.rawData = rawData
|
|
}
|
|
}
|
|
|
|
public enum VertexAIUsageFetcher {
|
|
private static let log = CodexBarLog.logger(LogCategories.vertexAIFetcher)
|
|
|
|
// Cloud Monitoring API endpoint for time series
|
|
private static let monitoringEndpoint = "https://monitoring.googleapis.com/v3/projects"
|
|
private static let usageWindowSeconds: TimeInterval = 24 * 60 * 60
|
|
|
|
public static func fetchUsage(
|
|
accessToken: String,
|
|
projectId: String?) async throws -> VertexAIUsageResponse
|
|
{
|
|
guard let projectId, !projectId.isEmpty else {
|
|
throw VertexAIFetchError.noProject
|
|
}
|
|
|
|
return try await Self.fetchQuotaUsage(
|
|
accessToken: accessToken,
|
|
projectId: projectId)
|
|
}
|
|
|
|
private static func fetchQuotaUsage(
|
|
accessToken: String,
|
|
projectId: String) async throws -> VertexAIUsageResponse
|
|
{
|
|
let usageFilter = """
|
|
metric.type="serviceruntime.googleapis.com/quota/allocation/usage" \
|
|
AND resource.type="consumer_quota" \
|
|
AND resource.label.service="aiplatform.googleapis.com"
|
|
"""
|
|
let limitFilter = """
|
|
metric.type="serviceruntime.googleapis.com/quota/limit" \
|
|
AND resource.type="consumer_quota" \
|
|
AND resource.label.service="aiplatform.googleapis.com"
|
|
"""
|
|
|
|
let usageSeries = try await Self.fetchTimeSeries(
|
|
accessToken: accessToken,
|
|
projectId: projectId,
|
|
filter: usageFilter)
|
|
let limitSeries = try await Self.fetchTimeSeries(
|
|
accessToken: accessToken,
|
|
projectId: projectId,
|
|
filter: limitFilter)
|
|
|
|
let usageByKey = Self.aggregate(series: usageSeries)
|
|
let limitByKey = Self.aggregate(series: limitSeries)
|
|
|
|
guard !usageByKey.isEmpty, !limitByKey.isEmpty else {
|
|
throw VertexAIFetchError.noData
|
|
}
|
|
|
|
var maxPercent: Double?
|
|
var matchedCount = 0
|
|
var matchedKeys: Set<QuotaKey> = []
|
|
for (key, limit) in limitByKey {
|
|
guard limit > 0, let usage = usageByKey[key] else { continue }
|
|
matchedKeys.insert(key)
|
|
matchedCount += 1
|
|
let percent = (usage / limit) * 100.0
|
|
maxPercent = max(maxPercent ?? percent, percent)
|
|
}
|
|
|
|
guard let usedPercent = maxPercent, matchedCount > 0 else {
|
|
throw VertexAIFetchError.noData
|
|
}
|
|
|
|
let unmatchedUsage = Set(usageByKey.keys).subtracting(matchedKeys).count
|
|
let unmatchedLimit = Set(limitByKey.keys).subtracting(matchedKeys).count
|
|
Self.log.debug("Quota series preview", metadata: [
|
|
"usageKeys": Self.previewKeys(usageByKey),
|
|
"limitKeys": Self.previewKeys(limitByKey),
|
|
])
|
|
Self.log.info("Parsed quota", metadata: [
|
|
"usedPercent": "\(usedPercent)",
|
|
"usageSeries": "\(usageByKey.count)",
|
|
"limitSeries": "\(limitByKey.count)",
|
|
"matchedSeries": "\(matchedCount)",
|
|
"unmatchedUsage": "\(unmatchedUsage)",
|
|
"unmatchedLimit": "\(unmatchedLimit)",
|
|
])
|
|
|
|
return VertexAIUsageResponse(
|
|
requestsUsedPercent: usedPercent,
|
|
tokensUsedPercent: nil,
|
|
resetsAt: nil,
|
|
resetDescription: nil,
|
|
rawData: nil)
|
|
}
|
|
|
|
private struct MonitoringTimeSeriesResponse: Decodable {
|
|
let timeSeries: [MonitoringTimeSeries]?
|
|
let nextPageToken: String?
|
|
}
|
|
|
|
private struct MonitoringTimeSeries: Decodable {
|
|
let metric: MonitoringMetric
|
|
let resource: MonitoringResource
|
|
let points: [MonitoringPoint]
|
|
}
|
|
|
|
private struct MonitoringMetric: Decodable {
|
|
let type: String?
|
|
let labels: [String: String]?
|
|
}
|
|
|
|
private struct MonitoringResource: Decodable {
|
|
let type: String?
|
|
let labels: [String: String]?
|
|
}
|
|
|
|
private struct MonitoringPoint: Decodable {
|
|
let value: MonitoringValue
|
|
}
|
|
|
|
private struct MonitoringValue: Decodable {
|
|
let doubleValue: Double?
|
|
let int64Value: String?
|
|
}
|
|
|
|
private struct QuotaKey: Hashable {
|
|
let quotaMetric: String
|
|
let limitName: String
|
|
let location: String
|
|
}
|
|
|
|
private static func fetchTimeSeries(
|
|
accessToken: String,
|
|
projectId: String,
|
|
filter: String) async throws -> [MonitoringTimeSeries]
|
|
{
|
|
let now = Date()
|
|
let start = now.addingTimeInterval(-Self.usageWindowSeconds)
|
|
let formatter = ISO8601DateFormatter()
|
|
var pageToken: String?
|
|
var allSeries: [MonitoringTimeSeries] = []
|
|
|
|
repeat {
|
|
guard var components = URLComponents(
|
|
string: "\(Self.monitoringEndpoint)/\(projectId)/timeSeries")
|
|
else {
|
|
throw VertexAIFetchError.invalidResponse("Invalid Monitoring URL")
|
|
}
|
|
|
|
var queryItems = [
|
|
URLQueryItem(name: "filter", value: filter),
|
|
URLQueryItem(name: "interval.startTime", value: formatter.string(from: start)),
|
|
URLQueryItem(name: "interval.endTime", value: formatter.string(from: now)),
|
|
URLQueryItem(name: "aggregation.alignmentPeriod", value: "3600s"),
|
|
URLQueryItem(name: "aggregation.perSeriesAligner", value: "ALIGN_MAX"),
|
|
URLQueryItem(name: "view", value: "FULL"),
|
|
]
|
|
if let pageToken {
|
|
queryItems.append(URLQueryItem(name: "pageToken", value: pageToken))
|
|
}
|
|
components.queryItems = queryItems
|
|
|
|
guard let url = components.url else {
|
|
throw VertexAIFetchError.invalidResponse("Invalid Monitoring URL")
|
|
}
|
|
|
|
var request = URLRequest(url: url)
|
|
request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization")
|
|
request.timeoutInterval = 30
|
|
|
|
let response: ProviderHTTPResponse
|
|
|
|
do {
|
|
response = try await ProviderHTTPClient.shared.response(for: request)
|
|
} catch {
|
|
throw VertexAIFetchError.networkError(error)
|
|
}
|
|
|
|
switch response.statusCode {
|
|
case 401:
|
|
throw VertexAIFetchError.unauthorized
|
|
case 403:
|
|
throw VertexAIFetchError.forbidden
|
|
case 200:
|
|
break
|
|
default:
|
|
let body = String(data: response.data, encoding: .utf8) ?? ""
|
|
throw VertexAIFetchError.invalidResponse("HTTP \(response.statusCode): \(body)")
|
|
}
|
|
|
|
let decoded = try JSONDecoder().decode(MonitoringTimeSeriesResponse.self, from: response.data)
|
|
if let series = decoded.timeSeries {
|
|
allSeries.append(contentsOf: series)
|
|
}
|
|
pageToken = decoded.nextPageToken?.isEmpty == false ? decoded.nextPageToken : nil
|
|
} while pageToken != nil
|
|
|
|
return allSeries
|
|
}
|
|
|
|
private static func aggregate(series: [MonitoringTimeSeries]) -> [QuotaKey: Double] {
|
|
var buckets: [QuotaKey: Double] = [:]
|
|
|
|
for entry in series {
|
|
guard let key = Self.quotaKey(from: entry),
|
|
let value = Self.maxPointValue(from: entry.points)
|
|
else {
|
|
continue
|
|
}
|
|
buckets[key] = max(buckets[key] ?? 0, value)
|
|
}
|
|
|
|
return buckets
|
|
}
|
|
|
|
private static func quotaKey(from series: MonitoringTimeSeries) -> QuotaKey? {
|
|
let metricLabels = series.metric.labels ?? [:]
|
|
let resourceLabels = series.resource.labels ?? [:]
|
|
let quotaMetric = metricLabels["quota_metric"]
|
|
?? resourceLabels["quota_id"]
|
|
guard let quotaMetric, !quotaMetric.isEmpty else { return nil }
|
|
let limitName = metricLabels["limit_name"] ?? ""
|
|
let location = resourceLabels["location"] ?? "global"
|
|
return QuotaKey(quotaMetric: quotaMetric, limitName: limitName, location: location)
|
|
}
|
|
|
|
private static func maxPointValue(from points: [MonitoringPoint]) -> Double? {
|
|
points.compactMap(self.pointValue).max()
|
|
}
|
|
|
|
private static func pointValue(from point: MonitoringPoint) -> Double? {
|
|
if let doubleValue = point.value.doubleValue { return doubleValue }
|
|
if let int64Value = point.value.int64Value { return Double(int64Value) }
|
|
return nil
|
|
}
|
|
|
|
private static func previewKeys(_ map: [QuotaKey: Double], maxCount: Int = 3) -> String {
|
|
guard !map.isEmpty else { return "none" }
|
|
let keys = map.keys.prefix(maxCount).map { key in
|
|
"\(key.quotaMetric)|\(key.limitName)|\(key.location)"
|
|
}
|
|
let suffix = map.count > maxCount ? " +\(map.count - maxCount)" : ""
|
|
return keys.joined(separator: ", ") + suffix
|
|
}
|
|
}
|