Rocky
10bc457c72
Unify PiP recovery across startup and foreground transitions, surface actionable status, and add privacy-safe analytics plus the refreshed onboarding and account experience. Update documentation assets and advance the release build to 84.
556 lines
20 KiB
Swift
556 lines
20 KiB
Swift
// ManagedLLMClient.swift
|
|
// OSGKeyboard · Shared
|
|
//
|
|
// LLMClient implementation for scope-limited managed polish, AI and agent calls.
|
|
|
|
import Foundation
|
|
|
|
public struct ManagedLLMClient: LLMClient {
|
|
public enum Capability: String, Sendable {
|
|
case polish
|
|
case assistant = "ai"
|
|
case agent
|
|
|
|
var grantScope: ManagedGatewayCapability {
|
|
switch self {
|
|
case .polish: .polish
|
|
case .assistant: .assistant
|
|
case .agent: .agent
|
|
}
|
|
}
|
|
|
|
var defaultTaskKind: ManagedGatewayTaskKind {
|
|
switch self {
|
|
case .polish: .dictationPolish
|
|
case .assistant: .aiQuestion
|
|
case .agent: .agentPlanning
|
|
}
|
|
}
|
|
}
|
|
|
|
private struct Attempt {
|
|
let input: String
|
|
let context: String?
|
|
let timeout: TimeInterval?
|
|
let options: LLMGenerationOptions
|
|
let requestId: String
|
|
let forceRefresh: Bool
|
|
|
|
func forcingRefresh() -> Self {
|
|
Self(
|
|
input: input,
|
|
context: context,
|
|
timeout: timeout,
|
|
options: options,
|
|
requestId: requestId,
|
|
forceRefresh: true
|
|
)
|
|
}
|
|
}
|
|
|
|
public let capability: Capability
|
|
public let taskKind: ManagedGatewayTaskKind
|
|
public let requestPurpose: ManagedGatewayRequestPurpose?
|
|
public let requestTimeout: TimeInterval
|
|
|
|
private let baseURL: URL
|
|
private let grants: GatewayGrantCoordinator
|
|
private let session: URLSession
|
|
private let requestId: @Sendable () -> String
|
|
|
|
public init(
|
|
capability: Capability,
|
|
taskKind: ManagedGatewayTaskKind? = nil,
|
|
requestPurpose: ManagedGatewayRequestPurpose? = nil,
|
|
grants: GatewayGrantCoordinator,
|
|
baseURL: URL = GatewayGrantCoordinator.defaultBaseURL,
|
|
session: URLSession = .shared,
|
|
requestTimeout: TimeInterval = 15,
|
|
requestId: @escaping @Sendable () -> String = { UUID().uuidString }
|
|
) {
|
|
self.capability = capability
|
|
self.taskKind = taskKind ?? capability.defaultTaskKind
|
|
self.requestPurpose = requestPurpose
|
|
self.grants = grants
|
|
self.baseURL = baseURL
|
|
self.session = session
|
|
self.requestTimeout = requestTimeout
|
|
self.requestId = requestId
|
|
}
|
|
|
|
public func polish(
|
|
_ text: String,
|
|
systemPrompt: String,
|
|
timeout: TimeInterval?
|
|
) async throws -> String {
|
|
try await polish(
|
|
text,
|
|
systemPrompt: systemPrompt,
|
|
timeout: timeout,
|
|
options: .polishDefault
|
|
)
|
|
}
|
|
|
|
public func polish(
|
|
_ text: String,
|
|
systemPrompt: String,
|
|
timeout: TimeInterval?,
|
|
options: LLMGenerationOptions
|
|
) async throws -> String {
|
|
try await executeBuffered(
|
|
input: text,
|
|
context: systemPrompt.nilIfEmpty,
|
|
timeout: timeout,
|
|
options: options
|
|
)
|
|
}
|
|
|
|
public func complete(
|
|
messages: [LLMRequest.Message],
|
|
timeout: TimeInterval?,
|
|
options: LLMGenerationOptions
|
|
) async throws -> String {
|
|
let payload = Self.payload(from: messages)
|
|
return try await executeBuffered(
|
|
input: payload.input,
|
|
context: payload.context,
|
|
timeout: timeout,
|
|
options: options
|
|
)
|
|
}
|
|
|
|
public func completeStreaming(
|
|
messages: [LLMRequest.Message],
|
|
timeout: TimeInterval?,
|
|
options: LLMGenerationOptions
|
|
) -> AsyncThrowingStream<LLMStreamEvent, Error> {
|
|
AsyncThrowingStream { continuation in
|
|
let task = Task {
|
|
do {
|
|
let payload = Self.payload(from: messages)
|
|
let logicalRequestId = requestId()
|
|
let attempt = Attempt(
|
|
input: payload.input,
|
|
context: payload.context,
|
|
timeout: timeout,
|
|
options: options,
|
|
requestId: logicalRequestId,
|
|
forceRefresh: false
|
|
)
|
|
var emittedVisibleText = false
|
|
do {
|
|
try await streamAttempt(attempt) { chunk in
|
|
emittedVisibleText = true
|
|
continuation.yield(.delta(chunk))
|
|
}
|
|
} catch ManagedGatewayError.invalidGrant where !emittedVisibleText {
|
|
try await streamAttempt(attempt.forcingRefresh()) { chunk in
|
|
emittedVisibleText = true
|
|
continuation.yield(.delta(chunk))
|
|
}
|
|
}
|
|
continuation.finish()
|
|
} catch is CancellationError {
|
|
continuation.finish(throwing: LLMError.cancelled)
|
|
} catch let error as URLError where error.code == .cancelled {
|
|
continuation.finish(throwing: LLMError.cancelled)
|
|
} catch let error as URLError where error.code == .timedOut {
|
|
continuation.finish(throwing: ManagedGatewayError.timeout)
|
|
} catch let error as LLMError {
|
|
continuation.finish(throwing: error)
|
|
} catch let error as ManagedGatewayError {
|
|
continuation.finish(throwing: error)
|
|
} catch {
|
|
continuation.finish(throwing: LLMError.transport(String(describing: error)))
|
|
}
|
|
}
|
|
continuation.onTermination = { _ in task.cancel() }
|
|
}
|
|
}
|
|
|
|
private func executeBuffered(
|
|
input: String,
|
|
context: String?,
|
|
timeout: TimeInterval?,
|
|
options: LLMGenerationOptions
|
|
) async throws -> String {
|
|
let logicalRequestId = requestId()
|
|
let attempt = Attempt(
|
|
input: input,
|
|
context: context,
|
|
timeout: timeout,
|
|
options: options,
|
|
requestId: logicalRequestId,
|
|
forceRefresh: false
|
|
)
|
|
do {
|
|
return try await bufferedAttempt(attempt)
|
|
} catch ManagedGatewayError.invalidGrant {
|
|
do {
|
|
return try await bufferedAttempt(attempt.forcingRefresh())
|
|
} catch ManagedGatewayError.invalidGrant {
|
|
try? await grants.clearGrant()
|
|
throw ManagedGatewayError.invalidGrant
|
|
}
|
|
}
|
|
}
|
|
|
|
private func bufferedAttempt(_ attempt: Attempt) async throws -> String {
|
|
let request = try await makeRequest(
|
|
attempt,
|
|
stream: false
|
|
)
|
|
do {
|
|
let (data, response) = try await session.data(for: request)
|
|
guard let http = response as? HTTPURLResponse else {
|
|
throw LLMError.transport("non-HTTP response")
|
|
}
|
|
guard (200..<300).contains(http.statusCode) else {
|
|
throw ManagedGatewayHTTP.error(
|
|
data: data,
|
|
status: http.statusCode,
|
|
requestId: attempt.requestId
|
|
)
|
|
}
|
|
return try Self.responseText(from: data)
|
|
} catch is CancellationError {
|
|
throw LLMError.cancelled
|
|
} catch let error as URLError where error.code == .cancelled {
|
|
throw LLMError.cancelled
|
|
} catch let error as URLError where error.code == .timedOut {
|
|
throw ManagedGatewayError.timeout
|
|
} catch let error as ManagedGatewayError {
|
|
throw error
|
|
} catch let error as LLMError {
|
|
throw error
|
|
} catch {
|
|
throw LLMError.transport(String(describing: error))
|
|
}
|
|
}
|
|
|
|
private func streamAttempt(
|
|
_ attempt: Attempt,
|
|
onDelta: @escaping (String) -> Void
|
|
) async throws {
|
|
guard capability != .agent else {
|
|
// The server validates agent output as one structured response.
|
|
let text = try await bufferedAttempt(attempt)
|
|
if !text.isEmpty {
|
|
onDelta(text)
|
|
}
|
|
return
|
|
}
|
|
|
|
let request = try await makeRequest(
|
|
attempt,
|
|
stream: true
|
|
)
|
|
for try await payload in ManagedGatewayStreamTransport.payloads(
|
|
session: session,
|
|
request: request,
|
|
requestId: attempt.requestId
|
|
) {
|
|
try Task.checkCancellation()
|
|
if let error = Self.streamingError(from: payload, requestId: attempt.requestId) {
|
|
throw error
|
|
}
|
|
if let delta = Self.streamingDelta(from: payload), !delta.isEmpty {
|
|
onDelta(delta)
|
|
}
|
|
}
|
|
}
|
|
|
|
private func makeRequest(_ attempt: Attempt, stream: Bool) async throws -> URLRequest {
|
|
let trimmedInput = attempt.input.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
guard !trimmedInput.isEmpty, trimmedInput.count <= 32_000 else {
|
|
throw ManagedGatewayError.server(
|
|
code: "invalid_request",
|
|
status: 400,
|
|
requestId: attempt.requestId
|
|
)
|
|
}
|
|
let boundedContext = attempt.context?
|
|
.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
.nilIfEmpty
|
|
guard (boundedContext?.count ?? 0) <= 32_000 else {
|
|
throw ManagedGatewayError.server(
|
|
code: "invalid_request",
|
|
status: 400,
|
|
requestId: attempt.requestId
|
|
)
|
|
}
|
|
|
|
let token = try await grants.accessToken(
|
|
for: capability.grantScope,
|
|
forceRefresh: attempt.forceRefresh
|
|
)
|
|
let body = ManagedGatewayTextRequest(
|
|
input: trimmedInput,
|
|
context: boundedContext,
|
|
maxOutputTokens: min(max(attempt.options.maxTokens ?? 512, 1), 4_096),
|
|
temperature: min(max(attempt.options.temperature ?? 0.2, 0), 1),
|
|
stream: stream,
|
|
taskKind: taskKind,
|
|
requestPurpose: requestPurpose
|
|
)
|
|
|
|
var request = URLRequest(
|
|
url: baseURL.appending(path: "v1/gateway/llm/\(capability.rawValue)")
|
|
)
|
|
request.httpMethod = "POST"
|
|
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
|
|
request.setValue(stream ? "text/event-stream" : "application/json", forHTTPHeaderField: "Accept")
|
|
request.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
|
|
request.setValue(attempt.requestId, forHTTPHeaderField: "X-Request-ID")
|
|
request.timeoutInterval = attempt.timeout ?? requestTimeout
|
|
request.httpBody = try JSONEncoder().encode(body)
|
|
return request
|
|
}
|
|
|
|
static func payload(
|
|
from messages: [LLMRequest.Message]
|
|
) -> (input: String, context: String?) {
|
|
guard let inputIndex = messages.lastIndex(where: { $0.role == "user" }) else {
|
|
return ("", Self.contextText(from: messages))
|
|
}
|
|
let input = messages[inputIndex].content
|
|
var contextMessages = messages
|
|
contextMessages.remove(at: inputIndex)
|
|
return (input, Self.contextText(from: contextMessages))
|
|
}
|
|
|
|
private static func contextText(from messages: [LLMRequest.Message]) -> String? {
|
|
messages
|
|
.filter { !$0.content.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty }
|
|
.map { "\($0.role):\n\($0.content)" }
|
|
.joined(separator: "\n\n")
|
|
.nilIfEmpty
|
|
}
|
|
|
|
static func responseText(from data: Data) throws -> String {
|
|
let raw = String(data: data, encoding: .utf8)?
|
|
.trimmingCharacters(in: .whitespacesAndNewlines) ?? ""
|
|
guard !raw.isEmpty else { throw LLMError.decoding("empty managed gateway response") }
|
|
guard let json = try? JSONSerialization.jsonObject(with: data) else {
|
|
return raw
|
|
}
|
|
return extractedText(from: json)?.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
.nilIfEmpty ?? raw
|
|
}
|
|
|
|
private static func extractedText(from value: Any) -> String? {
|
|
if let string = value as? String {
|
|
return string
|
|
}
|
|
if let array = value as? [Any] {
|
|
let values = array.compactMap(extractedText(from:))
|
|
return values.isEmpty ? nil : values.joined()
|
|
}
|
|
guard let object = value as? [String: Any] else { return nil }
|
|
|
|
for key in ["output_text", "text"] {
|
|
if let text = object[key] as? String, !text.isEmpty {
|
|
return text
|
|
}
|
|
}
|
|
if let content = object["content"] {
|
|
if let text = content as? String, !text.isEmpty {
|
|
return text
|
|
}
|
|
if let extracted = extractedText(from: content), !extracted.isEmpty {
|
|
return extracted
|
|
}
|
|
}
|
|
if let message = object["message"] as? [String: Any],
|
|
let extracted = extractedText(from: message) {
|
|
return extracted
|
|
}
|
|
if let choices = object["choices"] as? [[String: Any]],
|
|
let first = choices.first {
|
|
if let message = first["message"] as? [String: Any],
|
|
let extracted = extractedText(from: message) {
|
|
return extracted
|
|
}
|
|
if let text = first["text"] as? String {
|
|
return text
|
|
}
|
|
}
|
|
if let data = object["data"], let extracted = extractedText(from: data) {
|
|
return extracted
|
|
}
|
|
if let output = object["output"], let extracted = extractedText(from: output) {
|
|
return extracted
|
|
}
|
|
return nil
|
|
}
|
|
|
|
static func streamingDelta(from data: Data) -> String? {
|
|
if let value = LLMStreamDeltaParser.responsesOutputTextDelta(from: data)
|
|
?? LLMStreamDeltaParser.chatCompletionsDelta(from: data)
|
|
?? LLMStreamDeltaParser.anthropicTextDelta(from: data) {
|
|
return value
|
|
}
|
|
guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
|
|
return String(data: data, encoding: .utf8)
|
|
}
|
|
for key in ["delta", "text", "content"] {
|
|
if let value = json[key] as? String, !value.isEmpty {
|
|
return value
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
private static func streamingError(
|
|
from data: Data,
|
|
requestId: String
|
|
) -> ManagedGatewayError? {
|
|
guard let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
|
|
let code = object["code"] as? String,
|
|
object["message"] != nil || code.hasSuffix("_error") else {
|
|
return nil
|
|
}
|
|
if ["insufficient_credits", "insufficient_balance"].contains(code) {
|
|
return .insufficientCredits
|
|
}
|
|
if ["unauthorized", "gateway_grant_denied", "invalid_grant"].contains(code) {
|
|
return .invalidGrant
|
|
}
|
|
return .server(code: code, status: 200, requestId: requestId)
|
|
}
|
|
}
|
|
|
|
public enum ManagedGatewayLLMClientFactory {
|
|
public static func polish(
|
|
taskKind: ManagedGatewayTaskKind = .dictationPolish,
|
|
requestPurpose: ManagedGatewayRequestPurpose? = nil,
|
|
grants: GatewayGrantCoordinator,
|
|
baseURL: URL = GatewayGrantCoordinator.defaultBaseURL,
|
|
session: URLSession = .shared
|
|
) -> any LLMClient {
|
|
ManagedLLMClient(
|
|
capability: .polish,
|
|
taskKind: taskKind,
|
|
requestPurpose: requestPurpose,
|
|
grants: grants,
|
|
baseURL: baseURL,
|
|
session: session
|
|
)
|
|
}
|
|
|
|
public static func ai(
|
|
taskKind: ManagedGatewayTaskKind = .aiQuestion,
|
|
grants: GatewayGrantCoordinator,
|
|
baseURL: URL = GatewayGrantCoordinator.defaultBaseURL,
|
|
session: URLSession = .shared
|
|
) -> any LLMClient {
|
|
ManagedLLMClient(
|
|
capability: .assistant,
|
|
taskKind: taskKind,
|
|
grants: grants,
|
|
baseURL: baseURL,
|
|
session: session
|
|
)
|
|
}
|
|
|
|
public static func agent(
|
|
taskKind: ManagedGatewayTaskKind = .agentPlanning,
|
|
grants: GatewayGrantCoordinator,
|
|
baseURL: URL = GatewayGrantCoordinator.defaultBaseURL,
|
|
session: URLSession = .shared
|
|
) -> any LLMClient {
|
|
ManagedLLMClient(
|
|
capability: .agent,
|
|
taskKind: taskKind,
|
|
grants: grants,
|
|
baseURL: baseURL,
|
|
session: session
|
|
)
|
|
}
|
|
}
|
|
|
|
private extension String {
|
|
var nilIfEmpty: String? {
|
|
isEmpty ? nil : self
|
|
}
|
|
}
|
|
|
|
enum ManagedGatewayStreamTransport {
|
|
static func payloads(
|
|
session: URLSession,
|
|
request: URLRequest,
|
|
requestId: String
|
|
) -> AsyncThrowingStream<Data, Error> {
|
|
AsyncThrowingStream { continuation in
|
|
let task = Task {
|
|
do {
|
|
let (bytes, response) = try await session.bytes(for: request)
|
|
guard let http = response as? HTTPURLResponse else {
|
|
throw LLMError.transport("non-HTTP response")
|
|
}
|
|
guard (200..<300).contains(http.statusCode) else {
|
|
var body = Data()
|
|
for try await byte in bytes {
|
|
body.append(byte)
|
|
}
|
|
throw ManagedGatewayHTTP.error(
|
|
data: body,
|
|
status: http.statusCode,
|
|
requestId: requestId
|
|
)
|
|
}
|
|
|
|
let isEventStream = http.value(forHTTPHeaderField: "Content-Type")?
|
|
.lowercased()
|
|
.contains("text/event-stream") == true
|
|
if !isEventStream {
|
|
var body = Data()
|
|
for try await byte in bytes {
|
|
try Task.checkCancellation()
|
|
body.append(byte)
|
|
}
|
|
if !body.isEmpty {
|
|
continuation.yield(body)
|
|
}
|
|
continuation.finish()
|
|
return
|
|
}
|
|
|
|
var line = Data()
|
|
for try await byte in bytes {
|
|
try Task.checkCancellation()
|
|
if byte == UInt8(ascii: "\n") {
|
|
yieldSSELine(line, to: continuation)
|
|
line.removeAll(keepingCapacity: true)
|
|
} else if byte != UInt8(ascii: "\r") {
|
|
line.append(byte)
|
|
}
|
|
}
|
|
yieldSSELine(line, to: continuation)
|
|
continuation.finish()
|
|
} catch is CancellationError {
|
|
continuation.finish(throwing: LLMError.cancelled)
|
|
} catch let error as URLError where error.code == .cancelled {
|
|
continuation.finish(throwing: LLMError.cancelled)
|
|
} catch let error as URLError where error.code == .timedOut {
|
|
continuation.finish(throwing: ManagedGatewayError.timeout)
|
|
} catch {
|
|
continuation.finish(throwing: error)
|
|
}
|
|
}
|
|
continuation.onTermination = { _ in task.cancel() }
|
|
}
|
|
}
|
|
|
|
private static func yieldSSELine(
|
|
_ line: Data,
|
|
to continuation: AsyncThrowingStream<Data, Error>.Continuation
|
|
) {
|
|
guard let payload = LLMStreamTransport.sseDataPayload(fromLineBytes: line),
|
|
payload != Data("[DONE]".utf8) else {
|
|
return
|
|
}
|
|
continuation.yield(payload)
|
|
}
|
|
}
|