transui/macos/Sources/Transui/TranscriptionClient.swift

import Foundation

#if canImport(FoundationNetworking)
  import FoundationNetworking
#endif

enum TranscriptionError: LocalizedError {
  case invalidEndpoint
  case invalidResponse
  case http(Int)
  case allEndpointsFailed([String])

  var errorDescription: String? {
    switch self {
    case .invalidEndpoint: "The transcription endpoint is invalid."
    case .invalidResponse: "The server response did not contain a transcript."
    case .http(let status): "The server returned HTTP \(status)."
    case .allEndpointsFailed(let failures): failures.joined(separator: " ")
    }
  }
}

struct TranscriptionEndpoint: Equatable {
  let label: String
  let endpoint: String
  let model: String
  let key: String

  var isConfigured: Bool {
    !endpoint.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
  }
}

struct TranscriptionPlan {
  let hosted: TranscriptionEndpoint
  let local: TranscriptionEndpoint
  let hostedFirst: Bool

  var endpoints: [TranscriptionEndpoint] {
    var result: [TranscriptionEndpoint] = []
    if hostedFirst, hosted.isConfigured { result.append(hosted) }
    if local.isConfigured { result.append(local) }
    return result
  }
}

struct TranscriptionClient {
  static func endpointURL(from value: String) -> URL? {
    var endpoint = value.trimmingCharacters(in: .whitespacesAndNewlines)
    while endpoint.hasSuffix("/") { endpoint.removeLast() }
    if !endpoint.contains("/v1/audio/transcriptions") {
      endpoint += "/v1/audio/transcriptions"
    }
    guard let url = URL(string: endpoint),
      let scheme = url.scheme?.lowercased(),
      scheme == "http" || scheme == "https",
      url.host != nil
    else { return nil }
    return url
  }

  func transcribe(
    recording: URL,
    plan: TranscriptionPlan,
    prompt: String
  ) async throws -> String {
    let endpoints = plan.endpoints
    guard !endpoints.isEmpty else { throw TranscriptionError.invalidEndpoint }

    var failures: [String] = []
    for endpoint in endpoints {
      do {
        return try await transcribe(
          recording: recording,
          endpoint: endpoint.endpoint,
          model: endpoint.model,
          key: endpoint.key,
          prompt: prompt
        )
      } catch {
        failures.append("\(endpoint.label): \(error.localizedDescription)")
      }
    }
    throw TranscriptionError.allEndpointsFailed(failures)
  }

  private func transcribe(
    recording: URL,
    endpoint: String,
    model: String,
    key: String,
    prompt: String
  ) async throws -> String {
    guard let url = Self.endpointURL(from: endpoint) else {
      throw TranscriptionError.invalidEndpoint
    }

    let boundary = "transui-\(UUID().uuidString)"
    var request = URLRequest(url: url)
    request.httpMethod = "POST"
    request.timeoutInterval = 300
    request.setValue(
      "multipart/form-data; boundary=\(boundary)", forHTTPHeaderField: "Content-Type")
    request.setValue("application/json", forHTTPHeaderField: "Accept")
    if !key.isEmpty {
      request.setValue("Bearer \(key)", forHTTPHeaderField: "Authorization")
    }
    request.httpBody = try multipartBody(
      boundary: boundary,
      recording: recording,
      model: model,
      prompt: prompt
    )

    let (data, response) = try await URLSession.shared.data(for: request)
    guard let http = response as? HTTPURLResponse else {
      throw TranscriptionError.invalidResponse
    }
    guard (200..<300).contains(http.statusCode) else {
      throw TranscriptionError.http(http.statusCode)
    }
    guard data.count <= 1_048_576 else {
      throw TranscriptionError.invalidResponse
    }
    guard let object = try JSONSerialization.jsonObject(with: data) as? [String: Any],
      let text = object["text"] as? String
    else {
      throw TranscriptionError.invalidResponse
    }
    return TranscriptCleaner.clean(text)
  }

  private func multipartBody(
    boundary: String,
    recording: URL,
    model: String,
    prompt: String
  ) throws -> Data {
    var data = Data()
    func field(_ name: String, _ value: String) {
      data.append("--\(boundary)\r\n")
      data.append("Content-Disposition: form-data; name=\"\(name)\"\r\n\r\n")
      data.append("\(value)\r\n")
    }
    field("response_format", "json")
    if !model.isEmpty { field("model", model) }
    if !prompt.isEmpty { field("prompt", prompt) }
    data.append("--\(boundary)\r\n")
    data.append("Content-Disposition: form-data; name=\"file\"; filename=\"recording.mp3\"\r\n")
    data.append("Content-Type: audio/mpeg\r\n\r\n")
    data.append(try Data(contentsOf: recording))
    data.append("\r\n--\(boundary)--\r\n")
    return data
  }
}

extension Data {
  fileprivate mutating func append(_ string: String) {
    append(contentsOf: string.utf8)
  }
}