[Vertex AI] Bypass proxy for testing against staging

Author: andrewheard

FirebaseVertexAI/Sources/Constants.swift

@@ -17,7 +17,7 @@ import Foundation
 /// Constants associated with the Vertex AI for Firebase SDK.
 enum Constants {
   /// The Vertex AI backend endpoint URL.
-  static let baseURL = "https://firebasevertexai.googleapis.com"
+  static let baseURL = "staging-aiplatform.sandbox.googleapis.com"
 
   /// The base reverse-DNS name for `NSError` or `CustomNSError` error domains.
   ///

FirebaseVertexAI/Sources/CountTokensRequest.swift

@@ -17,6 +17,7 @@ import Foundation
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 struct CountTokensRequest {
   let model: String
+  let location: String
 
   let contents: [ModelContent]
   let systemInstruction: ModelContent?
@@ -31,7 +32,9 @@ extension CountTokensRequest: GenerativeAIRequest {
   typealias Response = CountTokensResponse
 
   var url: URL {
-    URL(string: "\(Constants.baseURL)/\(options.apiVersion)/\(model):countTokens")!
+    URL(
+      string: "https://\(location)-\(Constants.baseURL)/\(options.apiVersion)/\(model):countTokens"
+    )!
   }
 }
 

FirebaseVertexAI/Sources/GenerateContentRequest.swift

@@ -18,6 +18,7 @@ import Foundation
 struct GenerateContentRequest {
   /// Model name.
   let model: String
+  let location: String
   let contents: [ModelContent]
   let generationConfig: GenerationConfig?
   let safetySettings: [SafetySetting]?
@@ -45,7 +46,7 @@ extension GenerateContentRequest: GenerativeAIRequest {
   typealias Response = GenerateContentResponse
 
   var url: URL {
-    let modelURL = "\(Constants.baseURL)/\(options.apiVersion)/\(model)"
+    let modelURL = "https://\(location)-\(Constants.baseURL)/\(options.apiVersion)/\(model)"
     if isStreaming {
       return URL(string: "\(modelURL):streamGenerateContent?alt=sse")!
     } else {

FirebaseVertexAI/Sources/GenerativeAIRequest.swift

@@ -31,7 +31,7 @@ public struct RequestOptions {
   let timeout: TimeInterval
 
   /// The API version to use in requests to the backend.
-  let apiVersion = "v1beta"
+  let apiVersion = "v1beta1"
 
   /// Initializes a request options object.
   ///

FirebaseVertexAI/Sources/GenerativeAIService.swift

@@ -180,28 +180,20 @@ struct GenerativeAIService {
   private func urlRequest<T: GenerativeAIRequest>(request: T) async throws -> URLRequest {
     var urlRequest = URLRequest(url: request.url)
     urlRequest.httpMethod = "POST"
-    urlRequest.setValue(apiKey, forHTTPHeaderField: "x-goog-api-key")
+    guard let accessToken = ProcessInfo.processInfo.environment["GCLOUD_ACCESS_TOKEN"] else {
+      fatalError("""
+      Missing access token; run `gcloud auth print-access-token` and add an environment variable \
+      `GCLOUD_ACCESS_TOKEN` with the printed value.
+      Note: This value will only be valid for 60 minutes.
+      """)
+    }
+    urlRequest.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization")
     urlRequest.setValue(
       "\(GenerativeAIService.languageTag) \(GenerativeAIService.firebaseVersionTag)",
       forHTTPHeaderField: "x-goog-api-client"
     )
     urlRequest.setValue("application/json", forHTTPHeaderField: "Content-Type")
 
-    if let appCheck {
-      let tokenResult = await appCheck.getToken(forcingRefresh: false)
-      urlRequest.setValue(tokenResult.token, forHTTPHeaderField: "X-Firebase-AppCheck")
-      if let error = tokenResult.error {
-        VertexLog.error(
-          code: .appCheckTokenFetchFailed,
-          "Failed to fetch AppCheck token. Error: \(error)"
-        )
-      }
-    }
-
-    if let auth, let authToken = try await auth.getToken(forcingRefresh: false) {
-      urlRequest.setValue("Firebase \(authToken)", forHTTPHeaderField: "Authorization")
-    }
-
     let encoder = JSONEncoder()
     encoder.keyEncodingStrategy = .convertToSnakeCase
     urlRequest.httpBody = try encoder.encode(request)

FirebaseVertexAI/Sources/GenerativeModel.swift

@@ -26,6 +26,8 @@ public final class GenerativeModel {
   /// The backing service responsible for sending and receiving model requests to the backend.
   let generativeAIService: GenerativeAIService
 
+  let location: String
+
   /// Configuration parameters used for the MultiModalModel.
   let generationConfig: GenerationConfig?
 
@@ -61,6 +63,7 @@ public final class GenerativeModel {
   init(name: String,
        projectID: String,
        apiKey: String,
+       location: String = "us-central1",
        generationConfig: GenerationConfig? = nil,
        safetySettings: [SafetySetting]? = nil,
        tools: [Tool]?,
@@ -78,6 +81,7 @@ public final class GenerativeModel {
       auth: auth,
       urlSession: urlSession
     )
+    self.location = location
     self.generationConfig = generationConfig
     self.safetySettings = safetySettings
     self.tools = tools
@@ -125,6 +129,7 @@ public final class GenerativeModel {
     try content.throwIfError()
     let response: GenerateContentResponse
     let generateContentRequest = GenerateContentRequest(model: modelResourceName,
+                                                        location: location,
                                                         contents: content,
                                                         generationConfig: generationConfig,
                                                         safetySettings: safetySettings,
@@ -182,6 +187,7 @@ public final class GenerativeModel {
     -> AsyncThrowingStream<GenerateContentResponse, Error> {
     try content.throwIfError()
     let generateContentRequest = GenerateContentRequest(model: modelResourceName,
+                                                        location: location,
                                                         contents: content,
                                                         generationConfig: generationConfig,
                                                         safetySettings: safetySettings,
@@ -253,6 +259,7 @@ public final class GenerativeModel {
   public func countTokens(_ content: [ModelContent]) async throws -> CountTokensResponse {
     let countTokensRequest = CountTokensRequest(
       model: modelResourceName,
+      location: location,
       contents: content,
       systemInstruction: systemInstruction,
       tools: tools,

FirebaseVertexAI/Sources/VertexAI.swift

@@ -93,6 +93,7 @@ public class VertexAI {
       name: modelResourceName(modelName: modelName),
       projectID: projectID,
       apiKey: apiKey,
+      location: location,
       generationConfig: generationConfig,
       safetySettings: safetySettings,
       tools: tools,