[Firebase AI] Add support for thought summaries (#15096)

Author: andrewheard

FirebaseAI/CHANGELOG.md

@@ -1,5 +1,9 @@
+# 12.2.0
+- [feature] Added support for returning thought summaries, which are synthesized
+  versions of a model's internal reasoning process. (#15096)
+
 # 12.0.0
-- [added] Added support for Grounding with Google Search. (#15014)
+- [feature] Added support for Grounding with Google Search. (#15014)
 - [removed] Removed `CountTokensResponse.totalBillableCharacters` which was
   deprecated in 11.15.0. Use `totalTokens` instead. (#15056)
 

FirebaseAI/Sources/Chat.swift

@@ -147,31 +147,48 @@ public final class Chat: Sendable {
   }
 
   private func aggregatedChunks(_ chunks: [ModelContent]) -> ModelContent {
-    var parts: [any Part] = []
+    var parts: [InternalPart] = []
     var combinedText = ""
-    for aggregate in chunks {
-      // Loop through all the parts, aggregating the text and adding the images.
-      for part in aggregate.parts {
-        switch part {
-        case let textPart as TextPart:
-          combinedText += textPart.text
-
-        default:
-          // Don't combine it, just add to the content. If there's any text pending, add that as
-          // a part.
+    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 {
-            parts.append(TextPart(combinedText))
-            combinedText = ""
+            flush()
           }
-
-          parts.append(part)
+          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)
       }
     }
 
-    if !combinedText.isEmpty {
-      parts.append(TextPart(combinedText))
-    }
+    // Flush any remaining text.
+    flush()
 
     return ModelContent(role: "model", parts: parts)
   }

FirebaseAI/Sources/GenerateContentResponse.swift

@@ -57,30 +57,19 @@ public struct GenerateContentResponse: Sendable {
   public let usageMetadata: UsageMetadata?
 
   /// The response's content as text, if it exists.
+  ///
+  /// - Note: This does not include thought summaries; see ``thoughtSummary`` for more details.
   public var text: String? {
-    guard let candidate = candidates.first else {
-      AILog.error(
-        code: .generateContentResponseNoCandidates,
-        "Could not get text from a response that had no candidates."
-      )
-      return nil
-    }
-    let textValues: [String] = candidate.content.parts.compactMap { part in
-      switch part {
-      case let textPart as TextPart:
-        return textPart.text
-      default:
-        return nil
-      }
-    }
-    guard textValues.count > 0 else {
-      AILog.error(
-        code: .generateContentResponseNoText,
-        "Could not get a text part from the first candidate."
-      )
-      return nil
-    }
-    return textValues.joined(separator: " ")
+    return text(isThought: false)
+  }
+
+  /// A summary of the model's thinking process, if available.
+  ///
+  /// - Important: Thought summaries are only available when `includeThoughts` is enabled in the
+  ///   ``ThinkingConfig``. For more information, see the
+  ///   [Thinking](https://firebase.google.com/docs/ai-logic/thinking) documentation.
+  public var thoughtSummary: String? {
+    return text(isThought: true)
   }
 
   /// Returns function calls found in any `Part`s of the first candidate of the response, if any.
@@ -89,12 +78,10 @@ public struct GenerateContentResponse: Sendable {
       return []
     }
     return candidate.content.parts.compactMap { part in
-      switch part {
-      case let functionCallPart as FunctionCallPart:
-        return functionCallPart
-      default:
+      guard let functionCallPart = part as? FunctionCallPart, !part.isThought else {
         return nil
       }
+      return functionCallPart
     }
   }
 
@@ -107,7 +94,12 @@ public struct GenerateContentResponse: Sendable {
       """)
       return []
     }
-    return candidate.content.parts.compactMap { $0 as? InlineDataPart }
+    return candidate.content.parts.compactMap { part in
+      guard let inlineDataPart = part as? InlineDataPart, !part.isThought else {
+        return nil
+      }
+      return inlineDataPart
+    }
   }
 
   /// Initializer for SwiftUI previews or tests.
@@ -117,6 +109,30 @@ public struct GenerateContentResponse: Sendable {
     self.promptFeedback = promptFeedback
     self.usageMetadata = usageMetadata
   }
+
+  func text(isThought: Bool) -> String? {
+    guard let candidate = candidates.first else {
+      AILog.error(
+        code: .generateContentResponseNoCandidates,
+        "Could not get text from a response that had no candidates."
+      )
+      return nil
+    }
+    let textValues: [String] = candidate.content.parts.compactMap { part in
+      guard let textPart = part as? TextPart, part.isThought == isThought else {
+        return nil
+      }
+      return textPart.text
+    }
+    guard textValues.count > 0 else {
+      AILog.error(
+        code: .generateContentResponseNoText,
+        "Could not get a text part from the first candidate."
+      )
+      return nil
+    }
+    return textValues.joined(separator: " ")
+  }
 }
 
 /// A struct representing a possible reply to a content generation prompt. Each content generation

FirebaseAI/Sources/ModelContent.swift

@@ -31,41 +31,62 @@ extension [ModelContent] {
   }
 }
 
-/// A type describing data in media formats interpretable by an AI model. Each generative AI
-/// request or response contains an `Array` of ``ModelContent``s, and each ``ModelContent`` value
-/// may comprise multiple heterogeneous ``Part``s.
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
-public struct ModelContent: Equatable, Sendable {
-  enum InternalPart: Equatable, Sendable {
+struct InternalPart: Equatable, Sendable {
+  enum OneOfData: Equatable, Sendable {
     case text(String)
-    case inlineData(mimetype: String, Data)
-    case fileData(mimetype: String, uri: String)
+    case inlineData(InlineData)
+    case fileData(FileData)
     case functionCall(FunctionCall)
     case functionResponse(FunctionResponse)
   }
 
+  let data: OneOfData
+
+  let isThought: Bool?
+
+  let thoughtSignature: String?
+
+  init(_ data: OneOfData, isThought: Bool?, thoughtSignature: String?) {
+    self.data = data
+    self.isThought = isThought
+    self.thoughtSignature = thoughtSignature
+  }
+}
+
+/// A type describing data in media formats interpretable by an AI model. Each generative AI
+/// request or response contains an `Array` of ``ModelContent``s, and each ``ModelContent`` value
+/// may comprise multiple heterogeneous ``Part``s.
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+public struct ModelContent: Equatable, Sendable {
   /// The role of the entity creating the ``ModelContent``. For user-generated client requests,
   /// for example, the role is `user`.
   public let role: String?
 
   /// The data parts comprising this ``ModelContent`` value.
   public var parts: [any Part] {
-    var convertedParts = [any Part]()
-    for part in internalParts {
-      switch part {
+    return internalParts.map { part -> any Part in
+      switch part.data {
       case let .text(text):
-        convertedParts.append(TextPart(text))
-      case let .inlineData(mimetype, data):
-        convertedParts.append(InlineDataPart(data: data, mimeType: mimetype))
-      case let .fileData(mimetype, uri):
-        convertedParts.append(FileDataPart(uri: uri, mimeType: mimetype))
+        return TextPart(text, isThought: part.isThought, thoughtSignature: part.thoughtSignature)
+      case let .inlineData(inlineData):
+        return InlineDataPart(
+          inlineData, isThought: part.isThought, thoughtSignature: part.thoughtSignature
+        )
+      case let .fileData(fileData):
+        return FileDataPart(
+          fileData, isThought: part.isThought, thoughtSignature: part.thoughtSignature
+        )
       case let .functionCall(functionCall):
-        convertedParts.append(FunctionCallPart(functionCall))
+        return FunctionCallPart(
+          functionCall, isThought: part.isThought, thoughtSignature: part.thoughtSignature
+        )
       case let .functionResponse(functionResponse):
-        convertedParts.append(FunctionResponsePart(functionResponse))
+        return FunctionResponsePart(
+          functionResponse, isThought: part.isThought, thoughtSignature: part.thoughtSignature
+        )
       }
     }
-    return convertedParts
   }
 
   // TODO: Refactor this
@@ -78,17 +99,35 @@ public struct ModelContent: Equatable, Sendable {
     for part in parts {
       switch part {
       case let textPart as TextPart:
-        convertedParts.append(.text(textPart.text))
+        convertedParts.append(InternalPart(
+          .text(textPart.text),
+          isThought: textPart._isThought,
+          thoughtSignature: textPart.thoughtSignature
+        ))
       case let inlineDataPart as InlineDataPart:
-        let inlineData = inlineDataPart.inlineData
-        convertedParts.append(.inlineData(mimetype: inlineData.mimeType, inlineData.data))
+        convertedParts.append(InternalPart(
+          .inlineData(inlineDataPart.inlineData),
+          isThought: inlineDataPart._isThought,
+          thoughtSignature: inlineDataPart.thoughtSignature
+        ))
       case let fileDataPart as FileDataPart:
-        let fileData = fileDataPart.fileData
-        convertedParts.append(.fileData(mimetype: fileData.mimeType, uri: fileData.fileURI))
+        convertedParts.append(InternalPart(
+          .fileData(fileDataPart.fileData),
+          isThought: fileDataPart._isThought,
+          thoughtSignature: fileDataPart.thoughtSignature
+        ))
       case let functionCallPart as FunctionCallPart:
-        convertedParts.append(.functionCall(functionCallPart.functionCall))
+        convertedParts.append(InternalPart(
+          .functionCall(functionCallPart.functionCall),
+          isThought: functionCallPart._isThought,
+          thoughtSignature: functionCallPart.thoughtSignature
+        ))
       case let functionResponsePart as FunctionResponsePart:
-        convertedParts.append(.functionResponse(functionResponsePart.functionResponse))
+        convertedParts.append(InternalPart(
+          .functionResponse(functionResponsePart.functionResponse),
+          isThought: functionResponsePart._isThought,
+          thoughtSignature: functionResponsePart.thoughtSignature
+        ))
       default:
         fatalError()
       }
@@ -102,6 +141,11 @@ public struct ModelContent: Equatable, Sendable {
     let content = parts.flatMap { $0.partsValue }
     self.init(role: role, parts: content)
   }
+
+  init(role: String?, parts: [InternalPart]) {
+    self.role = role
+    internalParts = parts
+  }
 }
 
 // MARK: Codable Conformances
@@ -121,7 +165,29 @@ extension ModelContent: Codable {
 }
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
-extension ModelContent.InternalPart: Codable {
+extension InternalPart: Codable {
+  enum CodingKeys: String, CodingKey {
+    case isThought = "thought"
+    case thoughtSignature
+  }
+
+  public func encode(to encoder: Encoder) throws {
+    try data.encode(to: encoder)
+    var container = encoder.container(keyedBy: CodingKeys.self)
+    try container.encodeIfPresent(isThought, forKey: .isThought)
+    try container.encodeIfPresent(thoughtSignature, forKey: .thoughtSignature)
+  }
+
+  public init(from decoder: Decoder) throws {
+    data = try OneOfData(from: decoder)
+    let container = try decoder.container(keyedBy: CodingKeys.self)
+    isThought = try container.decodeIfPresent(Bool.self, forKey: .isThought)
+    thoughtSignature = try container.decodeIfPresent(String.self, forKey: .thoughtSignature)
+  }
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension InternalPart.OneOfData: Codable {
   enum CodingKeys: String, CodingKey {
     case text
     case inlineData
@@ -135,10 +201,10 @@ extension ModelContent.InternalPart: Codable {
     switch self {
     case let .text(text):
       try container.encode(text, forKey: .text)
-    case let .inlineData(mimetype, bytes):
-      try container.encode(InlineData(data: bytes, mimeType: mimetype), forKey: .inlineData)
-    case let .fileData(mimetype: mimetype, url):
-      try container.encode(FileData(fileURI: url, mimeType: mimetype), forKey: .fileData)
+    case let .inlineData(inlineData):
+      try container.encode(inlineData, forKey: .inlineData)
+    case let .fileData(fileData):
+      try container.encode(fileData, forKey: .fileData)
     case let .functionCall(functionCall):
       try container.encode(functionCall, forKey: .functionCall)
     case let .functionResponse(functionResponse):
@@ -151,11 +217,9 @@ extension ModelContent.InternalPart: Codable {
     if values.contains(.text) {
       self = try .text(values.decode(String.self, forKey: .text))
     } else if values.contains(.inlineData) {
-      let inlineData = try values.decode(InlineData.self, forKey: .inlineData)
-      self = .inlineData(mimetype: inlineData.mimeType, inlineData.data)
+      self = try .inlineData(values.decode(InlineData.self, forKey: .inlineData))
     } else if values.contains(.fileData) {
-      let fileData = try values.decode(FileData.self, forKey: .fileData)
-      self = .fileData(mimetype: fileData.mimeType, uri: fileData.fileURI)
+      self = try .fileData(values.decode(FileData.self, forKey: .fileData))
     } else if values.contains(.functionCall) {
       self = try .functionCall(values.decode(FunctionCall.self, forKey: .functionCall))
     } else if values.contains(.functionResponse) {

FirebaseAI/Sources/Types/Internal/InternalPart.swift

@@ -67,6 +67,9 @@ struct FunctionResponse: Codable, Equatable, Sendable {
 struct ErrorPart: Part, Error {
   let error: Error
 
+  let isThought = false
+  let thoughtSignature: String? = nil
+
   init(_ error: Error) {
     self.error = error
   }