firebase-ios-sdk/FirebaseAI/Sources/TemplateChat.swift

// 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

/// **[Public Preview]** An ongoing conversation with a model backed by a server prompt template.
///
/// > Warning: This API is a public preview and may be subject to change.
///
/// A `TemplateChat` manages conversation history on the client while utilizing the prompt,
/// system instructions, and configurations defined in a server prompt template. The template's
/// configurations apply to every turn automatically.
///
/// > Important: The referenced server prompt template must include the `{{history}}` tag in its
/// content, which instructs the template where to inject the conversation turns.
///
/// ### Function Calling in Chat
/// If the template defines function declarations in its frontmatter `tools`, the model may return
/// a ``FunctionCallPart`` in its response candidates. When this occurs, execute the local function
/// logic and supply the result as a ``FunctionResponsePart`` in a new message using one of the
/// `sendMessage(_:)` or `sendMessageStream(_:)` methods to continue the conversation.
///
/// For more details, see
/// [Multi-turn interactions with server prompt templates](https://firebase.google.com/docs/ai-logic/server-prompt-templates/multi-turn-interactions).
public final class TemplateChat: Sendable {
  private let model: TemplateGenerativeModel
  private let templateID: String
  /// The converted template inputs, or the error produced while converting them.
  ///
  /// Conversion happens eagerly in `init`, but `startChat` is non-throwing, so any failure is
  /// stored and rethrown from the `sendMessage`/`sendMessageStream` methods.
  private let inputs: Result<[String: TemplateInput], Error>
  private let _history: History

  init(model: TemplateGenerativeModel, templateID: String, inputs: [String: Any],
       history: [ModelContent]) {
    self.model = model
    self.templateID = templateID
    self.inputs = Result { try TemplateInput.inputs(from: inputs) }
    _history = History(history: history)
  }

  /// The previous content from the chat that has been successfully sent and received from the
  /// model.
  ///
  /// This history is provided to the model for each subsequent message sent, and is injected into
  /// the server prompt template at the position marked by the `{{history}}` tag.
  public var history: [ModelContent] {
    get {
      return _history.history
    }
    set {
      _history.history = newValue
    }
  }

  /// **[Public Preview]** Sends a message to the model and returns the response.
  ///
  /// > Warning: This API is a public preview and may be subject to change.
  ///
  /// - Parameter content: The message content to send to the model, including any
  ///   ``FunctionResponsePart``s responding to previous function calls; pass an empty array to
  ///   send an empty turn.
  /// - Returns: The content generated by the model.
  /// - Throws: A ``GenerateContentError`` if the request failed, or an `EncodingError` if `inputs`
  ///   cannot be encoded.
  public func sendMessage(_ content: [ModelContent]) async throws -> GenerateContentResponse {
    let newContent = content.map(populateContentRole)
    let response = try await model.generateContentWithHistory(
      history: _history.history + newContent,
      template: templateID,
      inputs: inputs.get()
    )
    _history.append(contentsOf: newContent)
    if let modelResponse = response.candidates.first {
      _history.append(modelResponse.content)
    }

    return response
  }

  /// **[Public Preview]** Sends a message to the model and returns the response.
  ///
  /// > Warning: This API is a public preview and may be subject to change.
  ///
  /// Calling this method with no arguments sends an empty turn, which runs the template's own
  /// prompt against the current history without adding any new content. This is typically how the
  /// first turn of a conversation is started when the template supplies the prompt itself:
  ///
  /// ```swift
  /// // Template content: "What was the weather in {{city}} on {{date}}?\n{{history}}"
  /// let chat = model.startChat(templateID: "weather", inputs: ["city": "Boston", ...])
  /// let response = try await chat.sendMessage()
  /// ```
  ///
  /// - Parameter parts: The new content to send as a single chat message; pass no arguments to
  ///   send an empty turn.
  /// - Returns: The content generated by the model.
  /// - Throws: A ``GenerateContentError`` if the request failed, or an `EncodingError` if `inputs`
  ///   cannot be encoded.
  public func sendMessage(_ parts: any PartsRepresentable...) async throws
    -> GenerateContentResponse {
    guard !parts.isEmpty else {
      return try await sendMessage([ModelContent]())
    }
    return try await sendMessage([ModelContent(parts: parts)])
  }

  /// **[Public Preview]** Sends a message to the model and returns the response as a stream of
  /// `GenerateContentResponse`s.
  ///
  /// > Warning: This API is a public preview and may be subject to change.
  ///
  /// - Parameter content: The message to send to the model, including any ``FunctionResponsePart``s
  ///   responding to previous function calls; pass an empty array to send an empty turn.
  /// - Returns: An `AsyncThrowingStream` that yields `GenerateContentResponse` objects.
  /// - Throws: A ``GenerateContentError`` if the request failed, or an `EncodingError` if `inputs`
  ///   cannot be encoded.
  @available(macOS 12.0, watchOS 8.0, *)
  public func sendMessageStream(_ content: [ModelContent]) throws
    -> AsyncThrowingStream<GenerateContentResponse, Error> {
    let newContent = content.map(populateContentRole)
    let stream = try model.generateContentStreamWithHistory(
      history: _history.history + newContent,
      template: templateID,
      inputs: inputs.get()
    )

    return AsyncThrowingStream { continuation in
      let task = 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 = _history.aggregatedChunks(aggregatedContent)
        _history.append(aggregated)
        continuation.finish()
      }

      continuation.onTermination = { _ in task.cancel() }
    }
  }

  /// **[Public Preview]** Sends a message to the model and returns the response as a stream of
  /// `GenerateContentResponse`s.
  ///
  /// > Warning: This API is a public preview and may be subject to change.
  ///
  /// Calling this method with no arguments sends an empty turn, which runs the template's own
  /// prompt against the current history without adding any new content.
  ///
  /// - Parameter parts: The new content to send as a single chat message; pass no arguments to
  ///   send an empty turn.
  /// - Returns: An `AsyncThrowingStream` that yields `GenerateContentResponse` objects.
  /// - Throws: A ``GenerateContentError`` if the request failed, or an `EncodingError` if `inputs`
  ///   cannot be encoded.
  @available(macOS 12.0, watchOS 8.0, *)
  public func sendMessageStream(_ parts: any PartsRepresentable...) throws
    -> AsyncThrowingStream<GenerateContentResponse, Error> {
    guard !parts.isEmpty else {
      return try sendMessageStream([ModelContent]())
    }
    return try sendMessageStream([ModelContent(parts: parts)])
  }

  // MARK: - Private

  private func populateContentRole(_ content: ModelContent) -> ModelContent {
    if content.role != nil {
      return content
    } else {
      return ModelContent(role: "user", parts: content.parts)
    }
  }
}