firebase-ios-sdk/FirebaseAI/Sources/GenerativeAIService.swift

// Copyright 2023 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import Foundation
import os.log

struct GenerativeAIService {
  let firebaseInfo: FirebaseInfo

  private let urlSession: URLSession

  init(firebaseInfo: FirebaseInfo, urlSession: URLSession) {
    self.firebaseInfo = firebaseInfo
    self.urlSession = urlSession
  }

  func loadRequest<T: GenerativeAIRequest>(request: T) async throws -> T.Response {
    let urlRequest = try await urlRequest(request: request)

    #if DEBUG
      if #available(macOS 11.0, *) {
        printCURLCommand(from: urlRequest)
      }
    #endif

    let data: Data
    let rawResponse: URLResponse
    (data, rawResponse) = try await urlSession.data(for: urlRequest)

    let response = try httpResponse(urlResponse: rawResponse)

    // Verify the status code is 200
    guard response.statusCode == 200 else {
      AILog.error(
        code: .loadRequestResponseError,
        "The server responded with an error: \(response)"
      )
      if let responseString = String(data: data, encoding: .utf8) {
        AILog.error(
          code: .loadRequestResponseErrorPayload,
          "Response payload: \(responseString)"
        )
      }

      throw parseError(responseData: data)
    }

    return try parseResponse(T.Response.self, from: data)
  }

  @available(macOS 12.0, watchOS 8.0, *)
  func loadRequestStream<T: GenerativeAIRequest>(request: T)
    -> AsyncThrowingStream<T.Response, Error> where T: Sendable {
    return AsyncThrowingStream { continuation in
      let task = Task {
        do {
          let urlRequest = try await self.urlRequest(request: request)

          #if DEBUG
            printCURLCommand(from: urlRequest)
          #endif

          let stream: URLSession.AsyncBytes
          let rawResponse: URLResponse
          (stream, rawResponse) = try await urlSession.bytes(for: urlRequest)

          let response = try httpResponse(urlResponse: rawResponse)

          // Verify the status code is 200
          guard response.statusCode == 200 else {
            AILog.error(
              code: .loadRequestStreamResponseError,
              "The server responded with an error: \(response)"
            )
            var responseBody = ""
            for try await line in stream.lines {
              responseBody += line + "\n"
            }

            AILog.error(
              code: .loadRequestStreamResponseErrorPayload,
              "Response payload: \(responseBody)"
            )
            continuation.finish(throwing: parseError(responseBody: responseBody))
            return
          }

          // Received lines that are not server-sent events (SSE); these are not prefixed with
          // "data:"
          var extraLines = ""

          for try await line in stream.lines {
            AILog.debug(code: .loadRequestStreamResponseLine, "Stream response: \(line)")

            if line.hasPrefix("data:") {
              // We can assume 5 characters since it's utf-8 encoded, removing `data:`.
              let jsonText = String(line.dropFirst(5))
              let data = try jsonData(jsonText: jsonText)
              let content = try parseResponse(T.Response.self, from: data)
              continuation.yield(content)
            } else {
              extraLines += line
            }
          }

          if extraLines.count > 0 {
            continuation.finish(throwing: parseError(responseBody: extraLines))
            return
          }

          continuation.finish(throwing: nil)
        } catch {
          continuation.finish(throwing: error)
        }
      }

      continuation.onTermination = { @Sendable _ in
        task.cancel()
      }
    }
  }

  // MARK: - Private Helpers

  private func urlRequest<T: GenerativeAIRequest>(request: T) async throws -> URLRequest {
    var urlRequest = try URLRequest(url: request.getURL())
    urlRequest.httpMethod = "POST"
    let additionalClientTags = TaskLocals.isHybridRequest ? ["hybrid"] : []
    try await firebaseInfo.applyHeaders(
      to: &urlRequest,
      additionalClientTags: additionalClientTags
    )

    let encoder = JSONEncoder()
    urlRequest.httpBody = try encoder.encode(request)
    urlRequest.timeoutInterval = request.options.timeout

    return urlRequest
  }

  private func httpResponse(urlResponse: URLResponse) throws -> HTTPURLResponse {
    // The following condition should always be true: "Whenever you make HTTP URL load requests, any
    // response objects you get back from the URLSession, NSURLConnection, or NSURLDownload class
    // are instances of the HTTPURLResponse class."
    guard let response = urlResponse as? HTTPURLResponse else {
      AILog.error(
        code: .generativeAIServiceNonHTTPResponse,
        "Response wasn't an HTTP response, internal error \(urlResponse)"
      )
      throw URLError(
        .badServerResponse,
        userInfo: [NSLocalizedDescriptionKey: "Response was not an HTTP response."]
      )
    }

    return response
  }

  private func jsonData(jsonText: String) throws -> Data {
    guard let data = jsonText.data(using: .utf8) else {
      throw DecodingError.dataCorrupted(DecodingError.Context(
        codingPath: [],
        debugDescription: "Could not parse response as UTF8."
      ))
    }
    return data
  }

  private func parseError(responseBody: String) -> Error {
    do {
      let data = try jsonData(jsonText: responseBody)
      return parseError(responseData: data)
    } catch {
      return error
    }
  }

  private func parseError(responseData: Data) -> Error {
    do {
      let rpcError = try JSONDecoder().decode(BackendError.self, from: responseData)
      logRPCError(rpcError)
      return rpcError
    } catch {
      let responseString = String(data: responseData, encoding: .utf8) ?? ""
      return UnrecognizedRPCError(responseBody: responseString)
    }
  }

  // Log specific RPC errors that cannot be mitigated or handled by user code.
  // These errors do not produce specific GenerateContentError or CountTokensError cases.
  private func logRPCError(_ error: BackendError) {
    let projectID = firebaseInfo.projectID
    if error.isFirebaseAILogicServiceDisabledError() {
      AILog.error(code: .agentPlatformInFirebaseAPIDisabled, """
      The Firebase AI Logic SDK requires the Firebase AI Logic API \
      (`firebasevertexai.googleapis.com`) to be enabled in your Firebase project. Enable this API \
      by visiting the Firebase Console at
      https://console.firebase.google.com/project/\(projectID)/genai/ and clicking "Get started". \
      If you enabled this API recently, wait a few minutes for the action to propagate to our \
      systems and then retry.
      """)
    }
  }

  private func parseResponse<T: Decodable>(_ type: T.Type, from data: Data) throws -> T {
    do {
      return try JSONDecoder().decode(type, from: data)
    } catch {
      if let json = String(data: data, encoding: .utf8) {
        AILog.error(code: .loadRequestParseResponseFailedJSON, "JSON response: \(json)")
      }
      AILog.error(
        code: .loadRequestParseResponseFailedJSONError,
        "Error decoding server JSON: \(error)"
      )
      throw error
    }
  }

  #if DEBUG
    @available(macOS 11.0, *)
    private func cURLCommand(from request: URLRequest) -> String {
      var returnValue = "curl "
      if let allHeaders = request.allHTTPHeaderFields {
        for (key, value) in allHeaders {
          returnValue += "-H '\(key): \(value)' "
        }
      }

      guard let url = request.url else { return "" }
      returnValue += "'\(url.absoluteString)' "

      guard let body = request.httpBody,
            let jsonStr = String(bytes: body, encoding: .utf8) else { return "" }
      let escapedJSON = jsonStr.replacingOccurrences(of: "'", with: "'\\''")
      returnValue += "-d '\(escapedJSON)'"

      return returnValue
    }

    @available(macOS 11.0, *)
    private func printCURLCommand(from request: URLRequest) {
      guard AILog.additionalLoggingEnabled() else {
        return
      }
      let command = cURLCommand(from: request)
      os_log(.debug, log: AILog.logObject, """
      \(AILog.service) Creating request with the equivalent cURL command:
      ----- cURL command -----
      \(command, privacy: .private)
      ------------------------
      """)
    }
  #endif // DEBUG
}