[AI] Server Prompt Templates (#15402)

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
Co-authored-by: Andrew Heard <[email protected]>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: 4 people

FirebaseAI/Sources/AILog.swift

@@ -87,6 +87,7 @@ enum AILog {
     case generateContentResponseEmptyCandidates = 4003
     case invalidWebsocketURL = 4004
     case duplicateLiveSessionSetupComplete = 4005
+    case malformedURL = 4006
 
     // SDK Debugging
     case loadRequestStreamResponseLine = 5000
@@ -138,6 +139,17 @@ enum AILog {
     log(level: .debug, code: code, message)
   }
 
+  @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+  static func makeInternalError(message: String, code: MessageCode) -> GenerateContentError {
+    let error = GenerateContentError.internalError(underlying: NSError(
+      domain: "\(Constants.baseErrorDomain).Internal",
+      code: code.rawValue,
+      userInfo: [NSLocalizedDescriptionKey: message]
+    ))
+    AILog.error(code: code, message)
+    return error
+  }
+
   /// Returns `true` if additional logging has been enabled via a launch argument.
   static func additionalLoggingEnabled() -> Bool {
     return ProcessInfo.processInfo.arguments.contains(enableArgumentKey)

FirebaseAI/Sources/Chat.swift

@@ -19,35 +19,21 @@ import Foundation
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 public final class Chat: Sendable {
   private let model: GenerativeModel
+  private let _history: History
 
-  /// Initializes a new chat representing a 1:1 conversation between model and user.
   init(model: GenerativeModel, history: [ModelContent]) {
     self.model = model
-    self.history = history
+    _history = History(history: history)
   }
 
-  private let historyLock = NSLock()
-  private nonisolated(unsafe) var _history: [ModelContent] = []
   /// The previous content from the chat that has been successfully sent and received from the
   /// model. This will be provided to the model for each message sent as context for the discussion.
   public var history: [ModelContent] {
     get {
-      historyLock.withLock { _history }
+      return _history.history
     }
     set {
-      historyLock.withLock { _history = newValue }
-    }
-  }
-
-  private func appendHistory(contentsOf: [ModelContent]) {
-    historyLock.withLock {
-      _history.append(contentsOf: contentsOf)
-    }
-  }
-
-  private func appendHistory(_ newElement: ModelContent) {
-    historyLock.withLock {
-      _history.append(newElement)
+      _history.history = newValue
     }
   }
 
@@ -87,8 +73,8 @@ public final class Chat: Sendable {
     let toAdd = ModelContent(role: "model", parts: reply.parts)
 
     // Append the request and successful result to history, then return the value.
-    appendHistory(contentsOf: newContent)
-    appendHistory(toAdd)
+    _history.append(contentsOf: newContent)
+    _history.append(toAdd)
     return result
   }
 
@@ -136,63 +122,16 @@ public final class Chat: Sendable {
         }
 
         // Save the request.
-        appendHistory(contentsOf: newContent)
+        _history.append(contentsOf: newContent)
 
         // Aggregate the content to add it to the history before we finish.
-        let aggregated = self.aggregatedChunks(aggregatedContent)
-        self.appendHistory(aggregated)
+        let aggregated = self._history.aggregatedChunks(aggregatedContent)
+        self._history.append(aggregated)
         continuation.finish()
       }
     }
   }
 
-  private func aggregatedChunks(_ chunks: [ModelContent]) -> ModelContent {
-    var parts: [InternalPart] = []
-    var combinedText = ""
-    var combinedThoughts = ""
-
-    func flush() {
-      if !combinedThoughts.isEmpty {
-        parts.append(InternalPart(.text(combinedThoughts), isThought: true, thoughtSignature: nil))
-        combinedThoughts = ""
-      }
-      if !combinedText.isEmpty {
-        parts.append(InternalPart(.text(combinedText), isThought: nil, thoughtSignature: nil))
-        combinedText = ""
-      }
-    }
-
-    // Loop through all the parts, aggregating the text.
-    for part in chunks.flatMap({ $0.internalParts }) {
-      // Only text parts may be combined.
-      if case let .text(text) = part.data, part.thoughtSignature == nil {
-        // Thought summaries must not be combined with regular text.
-        if part.isThought ?? false {
-          // If we were combining regular text, flush it before handling "thoughts".
-          if !combinedText.isEmpty {
-            flush()
-          }
-          combinedThoughts += text
-        } else {
-          // If we were combining "thoughts", flush it before handling regular text.
-          if !combinedThoughts.isEmpty {
-            flush()
-          }
-          combinedText += text
-        }
-      } else {
-        // This is a non-combinable part (not text), flush any pending text.
-        flush()
-        parts.append(part)
-      }
-    }
-
-    // Flush any remaining text.
-    flush()
-
-    return ModelContent(role: "model", parts: parts)
-  }
-
   /// 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

@@ -135,6 +135,28 @@ public final class FirebaseAI: Sendable {
     )
   }
 
+  /// Initializes a new `TemplateGenerativeModel`.
+  ///
+  /// - Returns: A new `TemplateGenerativeModel` instance.
+  public func templateGenerativeModel() -> TemplateGenerativeModel {
+    return TemplateGenerativeModel(
+      generativeAIService: GenerativeAIService(firebaseInfo: firebaseInfo,
+                                               urlSession: GenAIURLSession.default),
+      apiConfig: apiConfig
+    )
+  }
+
+  /// Initializes a new `TemplateImagenModel`.
+  ///
+  /// - Returns: A new `TemplateImagenModel` instance.
+  public func templateImagenModel() -> TemplateImagenModel {
+    return TemplateImagenModel(
+      generativeAIService: GenerativeAIService(firebaseInfo: firebaseInfo,
+                                               urlSession: GenAIURLSession.default),
+      apiConfig: apiConfig
+    )
+  }
+
   /// **[Public Preview]** Initializes a ``LiveGenerativeModel`` with the given parameters.
   ///
   /// - Note: Refer to [the Firebase docs on the Live

FirebaseAI/Sources/GenerateContentRequest.swift

@@ -73,15 +73,23 @@ extension GenerateContentRequest {
 extension GenerateContentRequest: GenerativeAIRequest {
   typealias Response = GenerateContentResponse
 
-  var url: URL {
+  func getURL() throws -> URL {
     let modelURL = "\(apiConfig.service.endpoint.rawValue)/\(apiConfig.version.rawValue)/\(model)"
+    let urlString: String
     switch apiMethod {
     case .generateContent:
-      return URL(string: "\(modelURL):\(apiMethod.rawValue)")!
+      urlString = "\(modelURL):\(apiMethod.rawValue)"
     case .streamGenerateContent:
-      return URL(string: "\(modelURL):\(apiMethod.rawValue)?alt=sse")!
+      urlString = "\(modelURL):\(apiMethod.rawValue)?alt=sse"
     case .countTokens:
-      fatalError("\(Self.self) should be a property of \(CountTokensRequest.self).")
+      throw AILog.makeInternalError(
+        message: "\(Self.self) should be a property of \(CountTokensRequest.self).",
+        code: .malformedURL
+      )
     }
+    guard let url = URL(string: urlString) else {
+      throw AILog.makeInternalError(message: "Malformed URL: \(urlString)", code: .malformedURL)
+    }
+    return url
   }
 }

FirebaseAI/Sources/GenerativeAIRequest.swift

@@ -18,7 +18,7 @@ import Foundation
 protocol GenerativeAIRequest: Sendable, Encodable {
   associatedtype Response: Sendable, Decodable
 
-  var url: URL { get }
+  func getURL() throws -> URL
 
   var options: RequestOptions { get }
 }

FirebaseAI/Sources/GenerativeAIService.swift

@@ -26,7 +26,7 @@ struct GenerativeAIService {
   /// The Firebase SDK version in the format `fire/<version>`.
   static let firebaseVersionTag = "fire/\(FirebaseVersion())"
 
-  private let firebaseInfo: FirebaseInfo
+  let firebaseInfo: FirebaseInfo
 
   private let urlSession: URLSession
 
@@ -167,7 +167,7 @@ struct GenerativeAIService {
   // MARK: - Private Helpers
 
   private func urlRequest<T: GenerativeAIRequest>(request: T) async throws -> URLRequest {
-    var urlRequest = URLRequest(url: request.url)
+    var urlRequest = try URLRequest(url: request.getURL())
     urlRequest.httpMethod = "POST"
     urlRequest.setValue(firebaseInfo.apiKey, forHTTPHeaderField: "x-goog-api-key")
     urlRequest.setValue(

FirebaseAI/Sources/History.swift

@@ -0,0 +1,94 @@
+// Copyright 2025 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.
+
+import Foundation
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+final class History: Sendable {
+  private let historyLock = NSLock()
+  private nonisolated(unsafe) var _history: [ModelContent] = []
+  /// The previous content from the chat that has been successfully sent and received from the
+  /// model. This will be provided to the model for each message sent as context for the discussion.
+  public var history: [ModelContent] {
+    get {
+      historyLock.withLock { _history }
+    }
+    set {
+      historyLock.withLock { _history = newValue }
+    }
+  }
+
+  init(history: [ModelContent]) {
+    self.history = history
+  }
+
+  func append(contentsOf: [ModelContent]) {
+    historyLock.withLock {
+      _history.append(contentsOf: contentsOf)
+    }
+  }
+
+  func append(_ newElement: ModelContent) {
+    historyLock.withLock {
+      _history.append(newElement)
+    }
+  }
+
+  func aggregatedChunks(_ chunks: [ModelContent]) -> ModelContent {
+    var parts: [InternalPart] = []
+    var combinedText = ""
+    var combinedThoughts = ""
+
+    func flush() {
+      if !combinedThoughts.isEmpty {
+        parts.append(InternalPart(.text(combinedThoughts), isThought: true, thoughtSignature: nil))
+        combinedThoughts = ""
+      }
+      if !combinedText.isEmpty {
+        parts.append(InternalPart(.text(combinedText), isThought: nil, thoughtSignature: nil))
+        combinedText = ""
+      }
+    }
+
+    // Loop through all the parts, aggregating the text.
+    for part in chunks.flatMap({ $0.internalParts }) {
+      // Only text parts may be combined.
+      if case let .text(text) = part.data, part.thoughtSignature == nil {
+        // Thought summaries must not be combined with regular text.
+        if part.isThought ?? false {
+          // If we were combining regular text, flush it before handling "thoughts".
+          if !combinedText.isEmpty {
+            flush()
+          }
+          combinedThoughts += text
+        } else {
+          // If we were combining "thoughts", flush it before handling regular text.
+          if !combinedThoughts.isEmpty {
+            flush()
+          }
+          combinedText += text
+        }
+      } else {
+        // This is a non-combinable part (not text), flush any pending text.
+        flush()
+        parts.append(part)
+      }
+    }
+
+    // Flush any remaining text.
+    flush()
+
+    return ModelContent(role: "model", parts: parts)
+  }
+}