fix(ai): Use location in websocket endpoint

FirebaseAI/Sources/FirebaseAI.swift

@@ -47,7 +47,6 @@
                                 useLimitedUseAppCheckTokens: Bool = false) -> FirebaseAI {
     let instance = createInstance(
       app: app,
-      location: backend.location,
       apiConfig: backend.apiConfig,
       useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
@@ -188,21 +187,19 @@
 
   let apiConfig: APIConfig
 
-  /// A map of active `FirebaseAI` instances keyed by the `FirebaseApp` name and the `location`,
-  /// in the format `appName:location`.
+  /// A map of active `FirebaseAI` instances keyed by the `FirebaseApp`,  the `APIConfig`, and
+  /// `useLimitedUseAppCheckTokens`.
   private nonisolated(unsafe) static var instances: [InstanceKey: FirebaseAI] = [:]
 
   /// Lock to manage access to the `instances` array to avoid race conditions.
   private nonisolated(unsafe) static var instancesLock: os_unfair_lock = .init()
 
-  let location: String?
-
   static let defaultVertexAIAPIConfig = APIConfig(
-    service: .vertexAI(endpoint: .firebaseProxyProd),
-    version: .v1beta
+    service: .vertexAI(endpoint: .firebaseProxyProd, location: "us-central1"),
+    version: .v1beta,
   )
 
-  static func createInstance(app: FirebaseApp?, location: String?,
+  static func createInstance(app: FirebaseApp?,
                              apiConfig: APIConfig,
                              useLimitedUseAppCheckTokens: Bool) -> FirebaseAI {
     guard let app = app ?? FirebaseApp.app() else {
@@ -216,7 +213,6 @@
 
     let instanceKey = InstanceKey(
       appName: app.name,
-      location: location,
       apiConfig: apiConfig,
       useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
@@ -225,15 +221,14 @@
     }
     let newInstance = FirebaseAI(
       app: app,
-      location: location,
       apiConfig: apiConfig,
       useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
     instances[instanceKey] = newInstance
     return newInstance
   }
 
-  init(app: FirebaseApp, location: String?, apiConfig: APIConfig,
+  init(app: FirebaseApp, apiConfig: APIConfig,
        useLimitedUseAppCheckTokens: Bool) {
     guard let projectID = app.options.projectID else {
       fatalError("The Firebase app named \"\(app.name)\" has no project ID in its configuration.")
@@ -254,7 +249,6 @@
       useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
     self.apiConfig = apiConfig
-    self.location = location
   }
 
   func modelResourceName(modelName: String) -> String {
@@ -268,17 +262,14 @@
     }
 
     switch apiConfig.service {
-    case .vertexAI:
-      return vertexAIModelResourceName(modelName: modelName)
+    case let .vertexAI(endpoint: _, location: location):
+      return vertexAIModelResourceName(modelName: modelName, location: location)
     case .googleAI:
       return developerModelResourceName(modelName: modelName)
     }
   }
 
-  private func vertexAIModelResourceName(modelName: String) -> String {
-    guard let location else {
-      fatalError("Location must be specified for the Firebase AI service.")
-    }
+  private func vertexAIModelResourceName(modelName: String, location: String) -> String {
     guard !location.isEmpty && location
       .allSatisfy({ !$0.isWhitespace && !$0.isNewline && $0 != "/" }) else {
       fatalError("""
@@ -307,7 +298,6 @@
   /// This type is `Hashable` so that it can be used as a key in the `instances` dictionary.
   private struct InstanceKey: Sendable, Hashable {
     let appName: String
-    let location: String?
     let apiConfig: APIConfig
     let useLimitedUseAppCheckTokens: Bool
   }

FirebaseAI/Sources/Types/Internal/APIConfig.swift

@@ -27,6 +27,7 @@ struct APIConfig: Sendable, Hashable, Encodable {
   /// - Parameters:
   ///   - service: The API service to use for generative AI.
   ///   - version: The version of the API to use.
+  ///   - location: The server location to use.
   init(service: Service, version: Version) {
     self.service = service
     self.version = version
@@ -45,7 +46,7 @@ extension APIConfig {
     /// See the [Cloud
     /// docs](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/inference) for
     /// more details.
-    case vertexAI(endpoint: Endpoint)
+    case vertexAI(endpoint: Endpoint, location: String)
 
     /// The Gemini Developer API provided by Google AI.
     ///
@@ -57,7 +58,7 @@ extension APIConfig {
     /// This must correspond with the API set in `service`.
     var endpoint: Endpoint {
       switch self {
-      case let .vertexAI(endpoint: endpoint):
+      case let .vertexAI(endpoint: endpoint, location: _):
         return endpoint
       case let .googleAI(endpoint: endpoint):
         return endpoint

FirebaseAI/Sources/Types/Internal/Live/LiveSessionService.swift

@@ -309,10 +309,9 @@ actor LiveSessionService {
   /// Will apply the required app check and auth headers, as the backend expects them.
   private nonisolated func createWebsocket() async throws -> AsyncWebSocket {
     let host = apiConfig.service.endpoint.rawValue.withoutPrefix("https://")
-    // TODO: (b/448722577) Set a location based on the api config
     let urlString = switch apiConfig.service {
-    case .vertexAI:
-      "wss://\(host)/ws/google.firebase.vertexai.v1beta.LlmBidiService/BidiGenerateContent/locations/us-central1"
+    case let .vertexAI(endpoint: _, location: location):
+      "wss://\(host)/ws/google.firebase.vertexai.v1beta.LlmBidiService/BidiGenerateContent/locations/\(location)"
     case .googleAI:
       "wss://\(host)/ws/google.firebase.vertexai.v1beta.GenerativeService/BidiGenerateContent"
     }

FirebaseAI/Sources/Types/Public/Backend.swift

@@ -25,26 +25,28 @@ public struct Backend {
   ///     for a list of supported locations.
   public static func vertexAI(location: String = "us-central1") -> Backend {
     return Backend(
-      apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseProxyProd), version: .v1beta),
-      location: location
+      apiConfig: APIConfig(
+        service: .vertexAI(endpoint: .firebaseProxyProd, location: location),
+        version: .v1beta,
+      )
     )
   }
 
   /// Initializes a `Backend` configured for the Google Developer API.
   public static func googleAI() -> Backend {
     return Backend(
-      apiConfig: APIConfig(service: .googleAI(endpoint: .firebaseProxyProd), version: .v1beta),
-      location: nil
+      apiConfig: APIConfig(
+        service: .googleAI(endpoint: .firebaseProxyProd),
+        version: .v1beta,
+      )
     )
   }
 
   // MARK: - Internal
 
   let apiConfig: APIConfig
-  let location: String?
 
-  init(apiConfig: APIConfig, location: String?) {
+  init(apiConfig: APIConfig) {
     self.apiConfig = apiConfig
-    self.location = location
   }
 }