[AI] Add hybrid inference support in GenerativeModelSession

FirebaseAI/Sources/Extensions/Internal/LanguageModelSession+ModelSession.swift

@@ -0,0 +1,174 @@
+// Copyright 2026 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.
+
+#if compiler(>=6.2.3) && canImport(FoundationModels)
+  import Foundation
+  import FoundationModels
+
+  @available(iOS 26.0, macOS 26.0, *)
+  @available(tvOS, unavailable)
+  @available(watchOS, unavailable)
+  extension FoundationModels.LanguageModelSession: _ModelSession {
+    public var _hasHistory: Bool {
+      if transcript.isEmpty {
+        return false
+      }
+
+      for entry in transcript {
+        switch entry {
+        case .instructions:
+          continue
+        case .prompt:
+          return true
+        case .toolCalls:
+          return true
+        case .toolOutput:
+          return true
+        case .response:
+          return true
+        @unknown default:
+          // Unknown entry type, assuming that it is session history.
+          return true
+        }
+      }
+
+      return false
+    }
+
+    public func _respond(to prompt: [any Part], schema: FirebaseAI.GenerationSchema?,
+                         includeSchemaInPrompt: Bool, options: GenerationConfig?) async throws
+      -> _ModelSessionResponse {
+      let prompt = try prompt.toFoundationModelsPrompt()
+
+      let response: FoundationModels.LanguageModelSession
+        .Response<FoundationModels.GeneratedContent>
+      if let schema {
+        response = try await respond(
+          to: prompt,
+          schema: schema.generationSchema,
+          includeSchemaInPrompt: includeSchemaInPrompt
+          // TODO: Add options: GenerationOptions
+        )
+      } else {
+        response = try await respond(
+          to: prompt,
+          schema: String.generationSchema
+          // TODO: Add options: GenerationOptions
+        )
+      }
+
+      // TODO: Extract common response handling code into a helper method.
+      let responseText: String
+      if schema == nil, case let .string(text) = response.rawContent.kind {
+        responseText = text
+      } else {
+        responseText = response.rawContent.jsonString
+      }
+
+      let generatedContent = response.rawContent.firebaseGeneratedContent
+      let modelContent = ModelContent(
+        role: "model",
+        parts: [InternalPart(.text(responseText), isThought: false, thoughtSignature: nil)]
+      )
+      let candidate = Candidate(
+        content: modelContent,
+        safetyRatings: [],
+        finishReason: nil,
+        citationMetadata: nil
+      )
+      let rawResponse = GenerateContentResponse(
+        candidates: [candidate],
+        modelVersion: FirebaseAI.SystemLanguageModel.modelName
+      )
+
+      return _ModelSessionResponse(rawContent: generatedContent, rawResponse: rawResponse)
+    }
+
+    public func _streamResponse(to prompt: [any Part],
+                                schema: FirebaseAI.GenerationSchema?,
+                                includeSchemaInPrompt: Bool,
+                                options: GenerationConfig?)
+      -> sending AsyncThrowingStream<_ModelSessionResponse, any Error> {
+      return AsyncThrowingStream { continuation in
+        let foundationModelsPrompt: Prompt
+        do {
+          foundationModelsPrompt = try prompt.toFoundationModelsPrompt()
+        } catch {
+          continuation.finish(throwing: error)
+          return
+        }
+
+        let stream: FoundationModels.LanguageModelSession
+          .ResponseStream<FoundationModels.GeneratedContent>
+        if let schema {
+          stream = streamResponse(
+            to: foundationModelsPrompt,
+            schema: schema.generationSchema,
+            includeSchemaInPrompt: includeSchemaInPrompt
+            // TODO: Add options: GenerationOptions
+          )
+        } else {
+          stream = streamResponse(
+            to: foundationModelsPrompt,
+            schema: String.generationSchema
+            // TODO: Check `includeSchemaInPrompt: includeSchemaInPrompt` behaviour with `String`
+            // TODO: Add options: GenerationOptions
+          )
+        }
+
+        let task = Task {
+          do {
+            for try await snapshot in stream {
+              // TODO: Extract common response handling code into a helper method.
+              let responseText: String
+              if schema == nil, case let .string(text) = snapshot.rawContent.kind {
+                responseText = text
+              } else {
+                responseText = snapshot.rawContent.jsonString
+              }
+
+              let generatedContent = snapshot.rawContent.firebaseGeneratedContent
+              let modelContent = ModelContent(
+                role: "model",
+                parts: [InternalPart(.text(responseText), isThought: false, thoughtSignature: nil)]
+              )
+              let candidate = Candidate(
+                content: modelContent,
+                safetyRatings: [],
+                finishReason: nil,
+                citationMetadata: nil
+              )
+              let rawResponse = GenerateContentResponse(
+                candidates: [candidate],
+                modelVersion: FirebaseAI.SystemLanguageModel.modelName
+              )
+
+              let response = _ModelSessionResponse(
+                rawContent: generatedContent,
+                rawResponse: rawResponse
+              )
+
+              continuation.yield(response)
+            }
+            continuation.finish()
+          } catch {
+            continuation.finish(throwing: error)
+            return
+          }
+        }
+        continuation.onTermination = { _ in task.cancel() }
+      }
+    }
+  }
+#endif // compiler(>=6.2.3) && canImport(FoundationModels)

FirebaseAI/Sources/Extensions/Internal/SystemLanguageModel+LanguageModel.swift

@@ -0,0 +1,77 @@
+// Copyright 2026 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.
+
+#if compiler(>=6.2.3) && canImport(FoundationModels)
+  import FoundationModels
+
+  extension FirebaseAI.SystemLanguageModel: LanguageModel {
+    static let modelName = "apple-foundation-models-system-language-model"
+
+    public var _modelName: String {
+      return FirebaseAI.SystemLanguageModel.modelName
+    }
+
+    public func _startSession(tools: [any ToolRepresentable]?,
+                              instructions: String?) throws -> any _ModelSession {
+      switch availability {
+      case .available:
+        break
+      case let .unavailable(reason):
+        throw GenerativeModelSession.GenerationError.assetsUnavailable(
+          GenerativeModelSession.GenerationError.Context(debugDescription: """
+          The Foundation Models `SystemLanguageModel` is unavailable: \(reason)
+          """)
+        )
+      }
+
+      #if canImport(FoundationModels) && IS_FOUNDATION_MODELS_SUPPORTED_PLATFORM
+        if #available(iOS 26.0, macOS 26.0, visionOS 26.0, *) {
+          var afmTools = [any FoundationModels.Tool]()
+          // Only function calling tools are supported by Foundation Models.
+          for tool in tools ?? [] {
+            // Skips any unsupported tools such as `GoogleMaps` or `CodeExecution` since they are
+            // only
+            // supported by Gemini models.
+            // TODO: Decide whether to throw for unsupported `FirebaseAILogic.Tool` types or ignore.
+            let functionDeclarations = tool.toolRepresentation.functionDeclarations ?? []
+            for functionDeclaration in functionDeclarations {
+              switch functionDeclaration.kind {
+              case .manual:
+                // TODO: Decide whether ignore manual function calling declarations, throw or assert.
+                continue
+              case let .foundationModels(afmTool):
+                guard let afmTool = afmTool as? (any FoundationModels.Tool) else {
+                  assertionFailure("""
+                  The function declaration "\(afmTool)" in the tool "\(tool)" is not a
+                  `FoundationModels.Tool` type.
+                  """)
+                  continue
+                }
+                afmTools.append(afmTool)
+              }
+            }
+          }
+          return LanguageModelSession(tools: afmTools, instructions: instructions)
+        }
+      #endif // canImport(FoundationModels) && IS_FOUNDATION_MODELS_SUPPORTED_PLATFORM
+
+      throw GenerativeModelSession.GenerationError.assetsUnavailable(
+        GenerativeModelSession.GenerationError.Context(debugDescription: """
+        Failed to start a `LanguageModelSession`. The Foundation Models `SystemLanguageModel` is not
+        available on the current platform.
+        """)
+      )
+    }
+  }
+#endif // compiler(>=6.2.3) && canImport(FoundationModels)

FirebaseAI/Sources/FirebaseAI.swift

@@ -108,6 +108,9 @@ public final class FirebaseAI: Sendable {
 
   // TODO: Remove the `#if compiler(>=6.2.3)` when Xcode 26.2 is the minimum supported version.
   #if compiler(>=6.2.3)
+
+    // TODO: Add public API for instantiating models to use with hybrid GenerativeModelSession.
+
     /// Creates a new `GenerativeModelSession` with the given model.
     ///
     /// - Important: **Public Preview** - This API is a public preview and may be subject to change.
@@ -122,14 +125,28 @@ public final class FirebaseAI: Sendable {
     ///   - instructions: System instructions that direct the model's behavior.
     public func generativeModelSession(model: String, tools: [any ToolRepresentable]? = nil,
                                        instructions: String? = nil) -> GenerativeModelSession {
-      let tools = tools?.map { $0.toolRepresentation }
-      let model = generativeModel(
-        modelName: model,
-        tools: tools,
-        systemInstruction: instructions.map { ModelContent(role: "system", parts: $0) }
+      let geminiModel = geminiModel(modelName: model)
+
+      return generativeModelSession(model: geminiModel, tools: tools, instructions: instructions)
+    }
+
+    // TODO: Update this testing API for hybrid GenerativeModelSession.
+    func geminiModel(modelName: String, safetySettings: [SafetySetting]? = nil,
+                     toolConfig: ToolConfig? = nil) -> any LanguageModel {
+      return GeminiModel(
+        modelName: modelName,
+        modelResourceName: modelResourceName(modelName: modelName),
+        firebaseInfo: firebaseInfo,
+        apiConfig: apiConfig,
+        safetySettings: safetySettings,
+        toolConfig: toolConfig
       )
+    }
 
-      return GenerativeModelSession(model: model)
+    // TODO: Update this testing API for hybrid GenerativeModelSession.
+    func generativeModelSession(model: any LanguageModel, tools: [any ToolRepresentable]? = nil,
+                                instructions: String? = nil) -> GenerativeModelSession {
+      return GenerativeModelSession(model: model, tools: tools, instructions: instructions)
     }
 
     #if canImport(FoundationModels)

FirebaseAI/Sources/GenerateContentResponse.swift

@@ -71,6 +71,8 @@ public struct GenerateContentResponse: Sendable {
 
   let responseID: String?
 
+  let modelVersion: String?
+
   /// The response's content as text, if it exists.
   ///
   /// - Note: This does not include thought summaries; see ``thoughtSummary`` for more details.
@@ -124,6 +126,17 @@ public struct GenerateContentResponse: Sendable {
     self.promptFeedback = promptFeedback
     self.usageMetadata = usageMetadata
     responseID = nil
+    modelVersion = nil
+  }
+
+  init(candidates: [Candidate], promptFeedback: PromptFeedback? = nil,
+       usageMetadata: UsageMetadata? = nil, responseID: String? = nil,
+       modelVersion: String? = nil) {
+    self.candidates = candidates
+    self.promptFeedback = promptFeedback
+    self.usageMetadata = usageMetadata
+    self.responseID = responseID
+    self.modelVersion = modelVersion
   }
 
   func text(isThought: Bool) -> String? {
@@ -448,6 +461,7 @@ extension GenerateContentResponse: Decodable {
     case promptFeedback
     case usageMetadata
     case responseID = "responseId"
+    case modelVersion
   }
 
   public init(from decoder: Decoder) throws {
@@ -474,6 +488,7 @@ extension GenerateContentResponse: Decodable {
     promptFeedback = try container.decodeIfPresent(PromptFeedback.self, forKey: .promptFeedback)
     usageMetadata = try container.decodeIfPresent(UsageMetadata.self, forKey: .usageMetadata)
     responseID = try container.decodeIfPresent(String.self, forKey: .responseID)
+    modelVersion = try container.decodeIfPresent(String.self, forKey: .modelVersion)
   }
 }