[AI] Refactor template models to reduce code duplication

FirebaseAI/Sources/Chat.swift

@@ -105,6 +105,33 @@ public final class Chat: Sendable {
     return result
   }
 
+  func sendMessage<T: GenerativeAIRequest>(_ content: [ModelContent],
+                                           request: T) async throws -> GenerateContentResponse
+    where T.Response == GenerateContentResponse {
+    // Ensure that the new content has the role set.
+    let newContent = content.map(populateContentRole(_:))
+
+    let result = try await GenerativeModel.generateContent(service: model.generativeAIService,
+                                                           request: request)
+
+    guard let reply = result.candidates.first?.content else {
+      let error = NSError(domain: "com.google.generative-ai",
+                          code: -1,
+                          userInfo: [
+                            NSLocalizedDescriptionKey: "No candidates with content available.",
+                          ])
+      throw GenerateContentError.internalError(underlying: error)
+    }
+
+    // Make sure we inject the role into the content received.
+    let toAdd = ModelContent(role: "model", parts: reply.parts)
+
+    // Append the request and successful result to history, then return the value.
+    _history.append(contentsOf: newContent)
+    _history.append(toAdd)
+    return result
+  }
+
   @available(macOS 12.0, watchOS 8.0, *)
   func sendMessageStream(_ content: [ModelContent], generationConfig: GenerationConfig?) throws
     -> AsyncThrowingStream<GenerateContentResponse, Error> {
@@ -145,6 +172,48 @@ public final class Chat: Sendable {
     }
   }
 
+  @available(macOS 12.0, watchOS 8.0, *)
+  func sendMessageStream<T: GenerativeAIRequest>(_ content: [ModelContent], request: T) throws
+    -> AsyncThrowingStream<GenerateContentResponse, Error>
+    where T.Response == GenerateContentResponse {
+    // Ensure that the new content has the role set.
+    let newContent: [ModelContent] = content.map(populateContentRole(_:))
+
+    let stream = try GenerativeModel.generateContentStream(
+      service: model.generativeAIService,
+      request: request
+    )
+    return AsyncThrowingStream { continuation in
+      Task {
+        var aggregatedContent: [ModelContent] = []
+
+        do {
+          for try await chunk in stream {
+            // Capture any content that's streaming. This should be populated if there's no error.
+            if let chunkContent = chunk.candidates.first?.content {
+              aggregatedContent.append(chunkContent)
+            }
+
+            // Pass along the chunk.
+            continuation.yield(chunk)
+          }
+        } catch {
+          // Rethrow the error that the underlying stream threw. Don't add anything to history.
+          continuation.finish(throwing: error)
+          return
+        }
+
+        // Save the request.
+        _history.append(contentsOf: newContent)
+
+        // Aggregate the content to add it to the history before we finish.
+        let aggregated = self._history.aggregatedChunks(aggregatedContent)
+        self._history.append(aggregated)
+        continuation.finish()
+      }
+    }
+  }
+
   /// Populates the `role` field with `user` if it doesn't exist. Required in chat sessions.
   private func populateContentRole(_ content: ModelContent) -> ModelContent {
     if content.role != nil {

FirebaseAI/Sources/FirebaseAI.swift

@@ -161,6 +161,7 @@ public final class FirebaseAI: Sendable {
   /// - Returns: A new `TemplateGenerativeModel` instance.
   public func templateGenerativeModel() -> TemplateGenerativeModel {
     return TemplateGenerativeModel(
+      generativeModel: generativeModel(modelName: "template"),
       generativeAIService: GenerativeAIService(firebaseInfo: firebaseInfo,
                                                urlSession: GenAIURLSession.default),
       apiConfig: apiConfig

FirebaseAI/Sources/GenerativeModel.swift

@@ -321,12 +321,10 @@ public final class GenerativeModel: Sendable {
 
   // MARK: - Internal
 
-  func generateContent(_ content: [ModelContent],
-                       generationConfig: GenerationConfig?) async throws
-    -> GenerateContentResponse {
+  func generateContentRequest(_ content: [ModelContent], generationConfig: GenerationConfig?) throws
+    -> GenerateContentRequest {
     try content.throwIfError()
-    let response: GenerateContentResponse
-    let generateContentRequest = GenerateContentRequest(
+    return GenerateContentRequest(
       model: modelResourceName,
       contents: content,
       generationConfig: generationConfig,
@@ -338,8 +336,25 @@ public final class GenerativeModel: Sendable {
       apiMethod: .generateContent,
       options: requestOptions
     )
+  }
+
+  func generateContent(_ content: [ModelContent], generationConfig: GenerationConfig?) async throws
+    -> GenerateContentResponse {
+    let generateContentRequest = try generateContentRequest(content,
+                                                            generationConfig: generationConfig)
+
+    return try await GenerativeModel.generateContent(
+      service: generativeAIService,
+      request: generateContentRequest
+    )
+  }
+
+  static func generateContent<T: GenerativeAIRequest>(service: GenerativeAIService,
+                                                      request: T) async throws
+    -> GenerateContentResponse where T.Response == GenerateContentResponse {
+    let response: GenerateContentResponse
     do {
-      response = try await generativeAIService.loadRequest(request: generateContentRequest)
+      response = try await service.loadRequest(request: request)
     } catch {
       throw GenerativeModel.generateContentError(from: error)
     }
@@ -382,8 +397,19 @@ public final class GenerativeModel: Sendable {
       options: requestOptions
     )
 
+    return try GenerativeModel.generateContentStream(
+      service: generativeAIService,
+      request: generateContentRequest
+    )
+  }
+
+  @available(macOS 12.0, watchOS 8.0, *)
+  static func generateContentStream<T: GenerativeAIRequest>(service: GenerativeAIService,
+                                                            request: T) throws
+    -> AsyncThrowingStream<GenerateContentResponse, Error>
+    where T.Response == GenerateContentResponse {
     return AsyncThrowingStream { continuation in
-      let responseStream = generativeAIService.loadRequestStream(request: generateContentRequest)
+      let responseStream = service.loadRequestStream(request: request)
       Task {
         do {
           var didYieldResponse = false

FirebaseAI/Sources/TemplateChatSession.swift

@@ -20,22 +20,25 @@
 ///
 /// **Public Preview**: This API is a public preview and may be subject to change.
 final class TemplateChatSession: Sendable {
-  private let model: TemplateGenerativeModel
+  private let templateModel: TemplateGenerativeModel
+  private let chat: Chat
   private let templateID: String
-  private let _history: History
+  private let inputs: [String: any Sendable]
 
-  init(model: TemplateGenerativeModel, templateID: String, history: [ModelContent]) {
-    self.model = model
+  init(model: TemplateGenerativeModel, templateID: String, inputs: [String: any Sendable],
+       history: [ModelContent]) {
+    templateModel = model
+    chat = model.generativeModel.startChat(history: history)
     self.templateID = templateID
-    _history = History(history: history)
+    self.inputs = inputs
   }
 
   public var history: [ModelContent] {
     get {
-      return _history.history
+      return chat.history
     }
     set {
-      _history.history = newValue
+      chat.history = newValue
     }
   }
 
@@ -45,28 +48,22 @@
   ///
   /// - Parameters:
   ///   - content: The message to send to the model.
-  ///   - inputs: A dictionary of variables to substitute into the template.
   ///   - options: The ``RequestOptions`` for the request, currently used to override default
   /// request timeout.
   /// - Returns: The content generated by the model.
   /// - Throws: A ``GenerateContentError`` if the request failed.
   func sendMessage(_ content: [ModelContent],
-                   inputs: [String: Any],
                    options: RequestOptions = RequestOptions()) async throws
     -> GenerateContentResponse {
-    let templateInputs = try inputs.mapValues { try TemplateInput(value: $0) }
-    let newContent = content.map(populateContentRole)
-    let response = try await model.generateContentWithHistory(
-      history: _history.history + newContent,
+    let request = try templateModel.templateGenerateContentRequest(
       template: templateID,
-      inputs: templateInputs,
+      inputs: inputs,
+      history: content,
+      stream: false,
       options: options
     )
-    _history.append(contentsOf: newContent)
-    if let modelResponse = response.candidates.first {
-      _history.append(modelResponse.content)
-    }
-    return response
+
+    return try await chat.sendMessage(content, request: request)
   }
 
   /// Sends a message to the model and returns the response.
@@ -75,18 +72,14 @@
   ///
   /// - Parameters:
   ///   - message: The message to send to the model.
-  ///   - inputs: A dictionary of variables to substitute into the template.
   ///   - options: The ``RequestOptions`` for the request, currently used to override default
   /// request timeout.
   /// - Returns: The content generated by the model.
   /// - Throws: A ``GenerateContentError`` if the request failed.
   func sendMessage(_ message: any PartsRepresentable,
-                   inputs: [String: Any],
                    options: RequestOptions = RequestOptions()) async throws
     -> GenerateContentResponse {
-    return try await sendMessage([ModelContent(parts: message.partsValue)],
-                                 inputs: inputs,
-                                 options: options)
+    return try await sendMessage([ModelContent(parts: message.partsValue)], options: options)
   }
 
   /// Sends a message to the model and returns the response as a stream of
@@ -103,46 +96,18 @@
   /// - Throws: A ``GenerateContentError`` if the request failed.
   @available(macOS 12.0, watchOS 8.0, *)
   func sendMessageStream(_ content: [ModelContent],
-                         inputs: [String: Any],
                          options: RequestOptions = RequestOptions()) throws
     -> AsyncThrowingStream<GenerateContentResponse, Error> {
-    let templateInputs = try inputs.mapValues { try TemplateInput(value: $0) }
-    let newContent = content.map(populateContentRole)
-    let stream = try model.generateContentStreamWithHistory(
-      history: _history.history + newContent,
+    let request = try templateModel.templateGenerateContentRequest(
       template: templateID,
-      inputs: templateInputs,
+      inputs: inputs,
+      history: content,
+      stream: true,
       options: options
     )
-    return AsyncThrowingStream { continuation in
-      Task {
-        var aggregatedContent: [ModelContent] = []
-
-        do {
-          for try await chunk in stream {
-            // Capture any content that's streaming. This should be populated if there's no error.
-            if let chunkContent = chunk.candidates.first?.content {
-              aggregatedContent.append(chunkContent)
-            }
-
-            // Pass along the chunk.
-            continuation.yield(chunk)
-          }
-        } catch {
-          // Rethrow the error that the underlying stream threw. Don't add anything to history.
-          continuation.finish(throwing: error)
-          return
-        }
-
-        // Save the request.
-        _history.append(contentsOf: newContent)
+    let templateInputs = try inputs.mapValues { try TemplateInput(value: $0) }
 
-        // Aggregate the content to add it to the history before we finish.
-        let aggregated = _history.aggregatedChunks(aggregatedContent)
-        _history.append(aggregated)
-        continuation.finish()
-      }
-    }
+    return try chat.sendMessageStream(content, request: request)
   }
 
   /// Sends a message to the model and returns the response as a stream of
@@ -159,19 +124,9 @@
   /// - Throws: A ``GenerateContentError`` if the request failed.
   @available(macOS 12.0, watchOS 8.0, *)
   func sendMessageStream(_ message: any PartsRepresentable,
-                         inputs: [String: Any],
                          options: RequestOptions = RequestOptions()) throws
     -> AsyncThrowingStream<GenerateContentResponse, Error> {
     return try sendMessageStream([ModelContent(parts: message.partsValue)],
-                                 inputs: inputs,
                                  options: options)
   }
-
-  private func populateContentRole(_ content: ModelContent) -> ModelContent {
-    if content.role != nil {
-      return content
-    } else {
-      return ModelContent(role: "user", parts: content.parts)
-    }
-  }
 }