fix(ai): Use location in websocket endpoint

FirebaseAI/Sources/FirebaseAI.swift

@@ -47,7 +47,6 @@ public final class FirebaseAI: Sendable {
                                 useLimitedUseAppCheckTokens: Bool = false) -> FirebaseAI {
     let instance = createInstance(
       app: app,
-      location: backend.location,
       apiConfig: backend.apiConfig,
       useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
@@ -188,21 +187,20 @@ public final class FirebaseAI: Sendable {
 
   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
+    version: .v1beta,
+    location: "us-central1"
   )
 
-  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 +214,6 @@ public final class FirebaseAI: Sendable {
 
     let instanceKey = InstanceKey(
       appName: app.name,
-      location: location,
       apiConfig: apiConfig,
       useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
@@ -225,15 +222,14 @@ public final class FirebaseAI: Sendable {
     }
     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 +250,6 @@ public final class FirebaseAI: Sendable {
       useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
     self.apiConfig = apiConfig
-    self.location = location
   }
 
   func modelResourceName(modelName: String) -> String {
@@ -276,7 +271,7 @@ public final class FirebaseAI: Sendable {
   }
 
   private func vertexAIModelResourceName(modelName: String) -> String {
-    guard let location else {
+    guard let location = apiConfig.location else {
       fatalError("Location must be specified for the Firebase AI service.")
     }
     guard !location.isEmpty && location
@@ -307,7 +302,6 @@ public final class FirebaseAI: Sendable {
   /// 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

@@ -22,14 +22,19 @@ struct APIConfig: Sendable, Hashable, Encodable {
   /// The version of the selected API to use, e.g., "v1".
   let version: Version
 
+  /// The server location to use, e.g., "us-central1"
+  let location: String?
+
   /// Initializes an API configuration.
   ///
   /// - Parameters:
   ///   - service: The API service to use for generative AI.
   ///   - version: The version of the API to use.
-  init(service: Service, version: Version) {
+  ///   - location: The server location to use.
+  init(service: Service, version: Version, location: String?) {
     self.service = service
     self.version = version
+    self.location = location
   }
 }
 

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"
+      "wss://\(host)/ws/google.firebase.vertexai.v1beta.LlmBidiService/BidiGenerateContent/locations/\(apiConfig.location ?? "us-central1")"
     case .googleAI:
       "wss://\(host)/ws/google.firebase.vertexai.v1beta.GenerativeService/BidiGenerateContent"
     }

FirebaseAI/Sources/Types/Public/Backend.swift

@@ -25,26 +25,30 @@ 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),
+        version: .v1beta,
+        location: location
+      )
     )
   }
 
   /// 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,
+        location: nil
+      )
     )
   }
 
   // MARK: - Internal
 
   let apiConfig: APIConfig
-  let location: String?
 
-  init(apiConfig: APIConfig, location: String?) {
+  init(apiConfig: APIConfig) {
     self.apiConfig = apiConfig
-    self.location = location
   }
 }