From 3c462827546c837bee2b2e401e671f6dc3fb6904 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Fri, 4 Jul 2025 11:00:49 -0400
Subject: [PATCH 01/20] Refactor `ModelContent.InternalPart` to use wrapper
 types

---
 FirebaseAI/Sources/ModelContent.swift | 32 ++++++++++++---------------
 1 file changed, 14 insertions(+), 18 deletions(-)

diff --git a/FirebaseAI/Sources/ModelContent.swift b/FirebaseAI/Sources/ModelContent.swift
index 7d82bd76445..517d086913c 100644
--- a/FirebaseAI/Sources/ModelContent.swift
+++ b/FirebaseAI/Sources/ModelContent.swift
@@ -38,8 +38,8 @@ extension [ModelContent] {
 public struct ModelContent: Equatable, Sendable {
   enum InternalPart: 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)
   }
@@ -55,10 +55,10 @@ public struct ModelContent: Equatable, Sendable {
       switch part {
       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))
+      case let .inlineData(inlineData):
+        convertedParts.append(InlineDataPart(inlineData))
+      case let .fileData(fileData):
+        convertedParts.append(FileDataPart(fileData))
       case let .functionCall(functionCall):
         convertedParts.append(FunctionCallPart(functionCall))
       case let .functionResponse(functionResponse):
@@ -80,11 +80,9 @@ public struct ModelContent: Equatable, Sendable {
       case let textPart as TextPart:
         convertedParts.append(.text(textPart.text))
       case let inlineDataPart as InlineDataPart:
-        let inlineData = inlineDataPart.inlineData
-        convertedParts.append(.inlineData(mimetype: inlineData.mimeType, inlineData.data))
+        convertedParts.append(.inlineData(inlineDataPart.inlineData))
       case let fileDataPart as FileDataPart:
-        let fileData = fileDataPart.fileData
-        convertedParts.append(.fileData(mimetype: fileData.mimeType, uri: fileData.fileURI))
+        convertedParts.append(.fileData(fileDataPart.fileData))
       case let functionCallPart as FunctionCallPart:
         convertedParts.append(.functionCall(functionCallPart.functionCall))
       case let functionResponsePart as FunctionResponsePart:
@@ -135,10 +133,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 +149,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) {

From 16687893a538bec1abca1e163903c937a605729c Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Fri, 4 Jul 2025 11:27:15 -0400
Subject: [PATCH 02/20] Move existing `InternalPart` contents into `OneOfData`

---
 FirebaseAI/Sources/ModelContent.swift | 48 ++++++++++++++++++++-------
 1 file changed, 36 insertions(+), 12 deletions(-)

diff --git a/FirebaseAI/Sources/ModelContent.swift b/FirebaseAI/Sources/ModelContent.swift
index 517d086913c..24738a2001c 100644
--- a/FirebaseAI/Sources/ModelContent.swift
+++ b/FirebaseAI/Sources/ModelContent.swift
@@ -31,12 +31,9 @@ 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(InlineData)
     case fileData(FileData)
@@ -44,6 +41,18 @@ public struct ModelContent: Equatable, Sendable {
     case functionResponse(FunctionResponse)
   }
 
+  let data: OneOfData
+
+  init(_ data: OneOfData) {
+    self.data = data
+  }
+}
+
+/// 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?
@@ -52,7 +61,7 @@ public struct ModelContent: Equatable, Sendable {
   public var parts: [any Part] {
     var convertedParts = [any Part]()
     for part in internalParts {
-      switch part {
+      switch part.data {
       case let .text(text):
         convertedParts.append(TextPart(text))
       case let .inlineData(inlineData):
@@ -78,15 +87,17 @@ 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)))
       case let inlineDataPart as InlineDataPart:
-        convertedParts.append(.inlineData(inlineDataPart.inlineData))
+        convertedParts.append(InternalPart(.inlineData(inlineDataPart.inlineData)))
       case let fileDataPart as FileDataPart:
-        convertedParts.append(.fileData(fileDataPart.fileData))
+        convertedParts.append(InternalPart(.fileData(fileDataPart.fileData)))
       case let functionCallPart as FunctionCallPart:
-        convertedParts.append(.functionCall(functionCallPart.functionCall))
+        convertedParts.append(InternalPart(.functionCall(functionCallPart.functionCall)))
       case let functionResponsePart as FunctionResponsePart:
-        convertedParts.append(.functionResponse(functionResponsePart.functionResponse))
+        convertedParts.append(
+          InternalPart(.functionResponse(functionResponsePart.functionResponse))
+        )
       default:
         fatalError()
       }
@@ -119,7 +130,20 @@ 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: CodingKey {}
+
+  public func encode(to encoder: Encoder) throws {
+    try data.encode(to: encoder)
+  }
+
+  public init(from decoder: Decoder) throws {
+    data = try OneOfData(from: decoder)
+  }
+}
+
+@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

From 6d3b0771f0548c48089c1ebee7292195a54d4e36 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Mon, 7 Jul 2025 14:51:24 -0400
Subject: [PATCH 03/20] Add `isThought` to `Part` protocol

---
 FirebaseAI/Sources/ModelContent.swift         | 44 ++++++++++++++-----
 .../Sources/Types/Internal/InternalPart.swift |  2 +
 FirebaseAI/Sources/Types/Public/Part.swift    | 44 +++++++++++++++----
 3 files changed, 69 insertions(+), 21 deletions(-)

diff --git a/FirebaseAI/Sources/ModelContent.swift b/FirebaseAI/Sources/ModelContent.swift
index 24738a2001c..7fdd7d428c0 100644
--- a/FirebaseAI/Sources/ModelContent.swift
+++ b/FirebaseAI/Sources/ModelContent.swift
@@ -43,8 +43,11 @@ struct InternalPart: Equatable, Sendable {
 
   let data: OneOfData
 
-  init(_ data: OneOfData) {
+  let isThought: Bool?
+
+  init(_ data: OneOfData, isThought: Bool?) {
     self.data = data
+    self.isThought = isThought
   }
 }
 
@@ -63,15 +66,15 @@ public struct ModelContent: Equatable, Sendable {
     for part in internalParts {
       switch part.data {
       case let .text(text):
-        convertedParts.append(TextPart(text))
+        convertedParts.append(TextPart(text, isThought: part.isThought))
       case let .inlineData(inlineData):
-        convertedParts.append(InlineDataPart(inlineData))
+        convertedParts.append(InlineDataPart(inlineData, isThought: part.isThought))
       case let .fileData(fileData):
-        convertedParts.append(FileDataPart(fileData))
+        convertedParts.append(FileDataPart(fileData, isThought: part.isThought))
       case let .functionCall(functionCall):
-        convertedParts.append(FunctionCallPart(functionCall))
+        convertedParts.append(FunctionCallPart(functionCall, isThought: part.isThought))
       case let .functionResponse(functionResponse):
-        convertedParts.append(FunctionResponsePart(functionResponse))
+        convertedParts.append(FunctionResponsePart(functionResponse, isThought: part.isThought))
       }
     }
     return convertedParts
@@ -87,16 +90,27 @@ public struct ModelContent: Equatable, Sendable {
     for part in parts {
       switch part {
       case let textPart as TextPart:
-        convertedParts.append(InternalPart(.text(textPart.text)))
+        convertedParts.append(InternalPart(.text(textPart.text), isThought: textPart._isThought))
       case let inlineDataPart as InlineDataPart:
-        convertedParts.append(InternalPart(.inlineData(inlineDataPart.inlineData)))
+        convertedParts.append(
+          InternalPart(.inlineData(inlineDataPart.inlineData), isThought: inlineDataPart._isThought)
+        )
       case let fileDataPart as FileDataPart:
-        convertedParts.append(InternalPart(.fileData(fileDataPart.fileData)))
+        convertedParts.append(
+          InternalPart(.fileData(fileDataPart.fileData), isThought: fileDataPart._isThought)
+        )
       case let functionCallPart as FunctionCallPart:
-        convertedParts.append(InternalPart(.functionCall(functionCallPart.functionCall)))
+        convertedParts.append(
+          InternalPart(
+            .functionCall(functionCallPart.functionCall), isThought: functionCallPart._isThought
+          )
+        )
       case let functionResponsePart as FunctionResponsePart:
         convertedParts.append(
-          InternalPart(.functionResponse(functionResponsePart.functionResponse))
+          InternalPart(
+            .functionResponse(functionResponsePart.functionResponse),
+            isThought: functionResponsePart._isThought
+          )
         )
       default:
         fatalError()
@@ -131,14 +145,20 @@ extension ModelContent: Codable {
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 extension InternalPart: Codable {
-  enum CodingKeys: CodingKey {}
+  enum CodingKeys: String, CodingKey {
+    case isThought = "thought"
+  }
 
   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)
   }
 
   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)
   }
 }
 
diff --git a/FirebaseAI/Sources/Types/Internal/InternalPart.swift b/FirebaseAI/Sources/Types/Internal/InternalPart.swift
index d543fb80f38..062ae8aa93d 100644
--- a/FirebaseAI/Sources/Types/Internal/InternalPart.swift
+++ b/FirebaseAI/Sources/Types/Internal/InternalPart.swift
@@ -67,6 +67,8 @@ struct FunctionResponse: Codable, Equatable, Sendable {
 struct ErrorPart: Part, Error {
   let error: Error
 
+  let isThought = false
+
   init(_ error: Error) {
     self.error = error
   }
diff --git a/FirebaseAI/Sources/Types/Public/Part.swift b/FirebaseAI/Sources/Types/Public/Part.swift
index 4890b725f4d..4fea490a7bf 100644
--- a/FirebaseAI/Sources/Types/Public/Part.swift
+++ b/FirebaseAI/Sources/Types/Public/Part.swift
@@ -18,7 +18,9 @@ import Foundation
 ///
 /// Within a single value of ``Part``, different data types may not mix.
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
-public protocol Part: PartsRepresentable, Codable, Sendable, Equatable {}
+public protocol Part: PartsRepresentable, Codable, Sendable, Equatable {
+  var isThought: Bool { get }
+}
 
 /// A text part containing a string value.
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
@@ -26,8 +28,17 @@ public struct TextPart: Part {
   /// Text value.
   public let text: String
 
+  public var isThought: Bool { _isThought ?? false }
+
+  let _isThought: Bool?
+
   public init(_ text: String) {
+    self.init(text, isThought: nil)
+  }
+
+  init(_ text: String, isThought: Bool?) {
     self.text = text
+    _isThought = isThought
   }
 }
 
@@ -45,6 +56,7 @@ public struct TextPart: Part {
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 public struct InlineDataPart: Part {
   let inlineData: InlineData
+  let _isThought: Bool?
 
   /// The data provided in the inline data part.
   public var data: Data { inlineData.data }
@@ -52,6 +64,8 @@ public struct InlineDataPart: Part {
   /// The IANA standard MIME type of the data.
   public var mimeType: String { inlineData.mimeType }
 
+  public var isThought: Bool { _isThought ?? false }
+
   /// Creates an inline data part from data and a MIME type.
   ///
   /// > Important: Supported input types depend on the model on the model being used; see [input
@@ -67,11 +81,12 @@ public struct InlineDataPart: Part {
   ///     requirements](https://firebase.google.com/docs/vertex-ai/input-file-requirements) for
   ///     supported values.
   public init(data: Data, mimeType: String) {
-    self.init(InlineData(data: data, mimeType: mimeType))
+    self.init(InlineData(data: data, mimeType: mimeType), isThought: nil)
   }
 
-  init(_ inlineData: InlineData) {
+  init(_ inlineData: InlineData, isThought: Bool?) {
     self.inlineData = inlineData
+    _isThought = isThought
   }
 }
 
@@ -79,9 +94,11 @@ public struct InlineDataPart: Part {
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 public struct FileDataPart: Part {
   let fileData: FileData
+  let _isThought: Bool?
 
   public var uri: String { fileData.fileURI }
   public var mimeType: String { fileData.mimeType }
+  public var isThought: Bool { _isThought ?? false }
 
   /// Constructs a new file data part.
   ///
@@ -93,11 +110,12 @@ public struct FileDataPart: Part {
   ///     requirements](https://firebase.google.com/docs/vertex-ai/input-file-requirements) for
   ///     supported values.
   public init(uri: String, mimeType: String) {
-    self.init(FileData(fileURI: uri, mimeType: mimeType))
+    self.init(FileData(fileURI: uri, mimeType: mimeType), isThought: nil)
   }
 
-  init(_ fileData: FileData) {
+  init(_ fileData: FileData, isThought: Bool?) {
     self.fileData = fileData
+    _isThought = isThought
   }
 }
 
@@ -105,6 +123,7 @@ public struct FileDataPart: Part {
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 public struct FunctionCallPart: Part {
   let functionCall: FunctionCall
+  let _isThought: Bool?
 
   /// The name of the function to call.
   public var name: String { functionCall.name }
@@ -112,6 +131,8 @@ public struct FunctionCallPart: Part {
   /// The function parameters and values.
   public var args: JSONObject { functionCall.args }
 
+  public var isThought: Bool { _isThought ?? false }
+
   /// Constructs a new function call part.
   ///
   /// > Note: A `FunctionCallPart` is typically received from the model, rather than created
@@ -121,11 +142,12 @@ public struct FunctionCallPart: Part {
   ///   - name: The name of the function to call.
   ///   - args: The function parameters and values.
   public init(name: String, args: JSONObject) {
-    self.init(FunctionCall(name: name, args: args))
+    self.init(FunctionCall(name: name, args: args), isThought: nil)
   }
 
-  init(_ functionCall: FunctionCall) {
+  init(_ functionCall: FunctionCall, isThought: Bool?) {
     self.functionCall = functionCall
+    _isThought = isThought
   }
 }
 
@@ -137,6 +159,7 @@ public struct FunctionCallPart: Part {
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 public struct FunctionResponsePart: Part {
   let functionResponse: FunctionResponse
+  let _isThought: Bool?
 
   /// The name of the function that was called.
   public var name: String { functionResponse.name }
@@ -144,16 +167,19 @@ public struct FunctionResponsePart: Part {
   /// The function's response or return value.
   public var response: JSONObject { functionResponse.response }
 
+  public var isThought: Bool { _isThought ?? false }
+
   /// Constructs a new `FunctionResponse`.
   ///
   /// - Parameters:
   ///   - name: The name of the function that was called.
   ///   - response: The function's response.
   public init(name: String, response: JSONObject) {
-    self.init(FunctionResponse(name: name, response: response))
+    self.init(FunctionResponse(name: name, response: response), isThought: nil)
   }
 
-  init(_ functionResponse: FunctionResponse) {
+  init(_ functionResponse: FunctionResponse, isThought: Bool?) {
     self.functionResponse = functionResponse
+    _isThought = isThought
   }
 }

From f122e834abec45804f17abd991aba7aba04a4ed3 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Mon, 7 Jul 2025 15:44:49 -0400
Subject: [PATCH 04/20] Add `includeThoughts` to `ThinkingConfig`

---
 FirebaseAI/Sources/Types/Public/ThinkingConfig.swift | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/FirebaseAI/Sources/Types/Public/ThinkingConfig.swift b/FirebaseAI/Sources/Types/Public/ThinkingConfig.swift
index c0e8f31465b..c9e2cc43f4e 100644
--- a/FirebaseAI/Sources/Types/Public/ThinkingConfig.swift
+++ b/FirebaseAI/Sources/Types/Public/ThinkingConfig.swift
@@ -37,12 +37,15 @@ public struct ThinkingConfig: Sendable {
   /// feature or if the specified budget is not within the model's supported range.
   let thinkingBudget: Int?
 
+  let includeThoughts: Bool?
+
   /// Initializes a new `ThinkingConfig`.
   ///
   /// - Parameters:
   ///   - thinkingBudget: The maximum number of tokens to be used for the model's thinking process.
-  public init(thinkingBudget: Int? = nil) {
+  public init(thinkingBudget: Int? = nil, includeThoughts: Bool? = nil) {
     self.thinkingBudget = thinkingBudget
+    self.includeThoughts = includeThoughts
   }
 }
 

From 58c2439a6fe6006482c2078b18ea2379d5b286a4 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Tue, 8 Jul 2025 19:01:23 -0400
Subject: [PATCH 05/20] Add `includeThoughts` checks in integration tests

---
 .../GenerateContentIntegrationTests.swift     | 72 ++++++++++++++-----
 1 file changed, 54 insertions(+), 18 deletions(-)

diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
index 5a63ca41a6d..1df37a23825 100644
--- a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
+++ b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
@@ -134,47 +134,83 @@ struct GenerateContentIntegrationTests {
     #expect(candidatesTokensDetails.tokenCount == usageMetadata.candidatesTokenCount)
   }
 
-  @Test(arguments: [
-    (InstanceConfig.vertexAI_v1beta, ModelNames.gemini2_5_Flash, 0),
-    (InstanceConfig.vertexAI_v1beta, ModelNames.gemini2_5_Flash, 24576),
-    (InstanceConfig.vertexAI_v1beta_global, ModelNames.gemini2_5_Pro, 128),
-    (InstanceConfig.vertexAI_v1beta_global, ModelNames.gemini2_5_Pro, 32768),
-    (InstanceConfig.googleAI_v1beta, ModelNames.gemini2_5_Flash, 0),
-    (InstanceConfig.googleAI_v1beta, ModelNames.gemini2_5_Flash, 24576),
-    (InstanceConfig.googleAI_v1beta, ModelNames.gemini2_5_Pro, 128),
-    (InstanceConfig.googleAI_v1beta, ModelNames.gemini2_5_Pro, 32768),
-    (InstanceConfig.googleAI_v1beta_freeTier, ModelNames.gemini2_5_Flash, 0),
-    (InstanceConfig.googleAI_v1beta_freeTier, ModelNames.gemini2_5_Flash, 24576),
-  ])
+  @Test(
+    arguments: [
+      (.vertexAI_v1beta, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: 0)),
+      (.vertexAI_v1beta, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: 24576)),
+      (.vertexAI_v1beta, ModelNames.gemini2_5_Flash, ThinkingConfig(
+        thinkingBudget: 24576, includeThoughts: true
+      )),
+      (.vertexAI_v1beta_global, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: 128)),
+      (.vertexAI_v1beta_global, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: 32768)),
+      (.vertexAI_v1beta_global, ModelNames.gemini2_5_Pro, ThinkingConfig(
+        thinkingBudget: 32768, includeThoughts: true
+      )),
+      (.googleAI_v1beta, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: 0)),
+      (.googleAI_v1beta, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: 24576)),
+      (.googleAI_v1beta, ModelNames.gemini2_5_Flash, ThinkingConfig(
+        thinkingBudget: 24576, includeThoughts: true
+      )),
+      (.googleAI_v1beta, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: 128)),
+      (.googleAI_v1beta, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: 32768)),
+      (.googleAI_v1beta, ModelNames.gemini2_5_Pro, ThinkingConfig(
+        thinkingBudget: 32768, includeThoughts: true
+      )),
+      (.googleAI_v1beta_freeTier, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: 0)),
+      (
+        .googleAI_v1beta_freeTier,
+        ModelNames.gemini2_5_Flash,
+        ThinkingConfig(thinkingBudget: 24576)
+      ),
+      (.googleAI_v1beta_freeTier, ModelNames.gemini2_5_Flash, ThinkingConfig(
+        thinkingBudget: 24576, includeThoughts: true
+      )),
+      (.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemini2_5_Flash, ThinkingConfig(
+        thinkingBudget: 0
+      )),
+      (.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemini2_5_Flash, ThinkingConfig(
+        thinkingBudget: 24576
+      )),
+      (.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemini2_5_Flash, ThinkingConfig(
+        thinkingBudget: 24576, includeThoughts: true
+      )),
+    ] as [(InstanceConfig, String, ThinkingConfig)]
+  )
   func generateContentThinking(_ config: InstanceConfig, modelName: String,
-                               thinkingBudget: Int) async throws {
+                               thinkingConfig: ThinkingConfig) async throws {
     let model = FirebaseAI.componentInstance(config).generativeModel(
       modelName: modelName,
       generationConfig: GenerationConfig(
         temperature: 0.0,
         topP: 0.0,
         topK: 1,
-        thinkingConfig: ThinkingConfig(thinkingBudget: thinkingBudget)
+        thinkingConfig: thinkingConfig
       ),
       safetySettings: safetySettings
     )
+    let chat = model.startChat()
     let prompt = "Where is Google headquarters located? Answer with the city name only."
 
-    let response = try await model.generateContent(prompt)
+    let response = try await chat.sendMessage(prompt)
 
     let text = try #require(response.text).trimmingCharacters(in: .whitespacesAndNewlines)
     #expect(text == "Mountain View")
 
+    let candidate = try #require(response.candidates.first)
+    let thoughtParts = candidate.content.parts.compactMap { $0.isThought ? $0 : nil }
+    #expect(thoughtParts.isEmpty != thinkingConfig.includeThoughts)
+
     let usageMetadata = try #require(response.usageMetadata)
     #expect(usageMetadata.promptTokenCount.isEqual(to: 13, accuracy: tokenCountAccuracy))
     #expect(usageMetadata.promptTokensDetails.count == 1)
     let promptTokensDetails = try #require(usageMetadata.promptTokensDetails.first)
     #expect(promptTokensDetails.modality == .text)
     #expect(promptTokensDetails.tokenCount == usageMetadata.promptTokenCount)
-    if thinkingBudget == 0 {
-      #expect(usageMetadata.thoughtsTokenCount == 0)
-    } else {
+    if let thinkingBudget = thinkingConfig.thinkingBudget, thinkingBudget > 0 {
+      #expect(usageMetadata.thoughtsTokenCount > 0)
       #expect(usageMetadata.thoughtsTokenCount <= thinkingBudget)
+    } else {
+      #expect(usageMetadata.thoughtsTokenCount == 0)
     }
     #expect(usageMetadata.candidatesTokenCount.isEqual(to: 3, accuracy: tokenCountAccuracy))
     // The `candidatesTokensDetails` field is erroneously omitted when using the Google AI (Gemini

From ab38945829c8edd8fb6823fbc08175e8e85cc499 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Tue, 8 Jul 2025 19:02:42 -0400
Subject: [PATCH 06/20] Filter thoughts out of all convenience accessors

---
 FirebaseAI/Sources/Chat.swift                 | 10 ++++++----
 .../Sources/GenerateContentResponse.swift     | 19 ++++++++++---------
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/FirebaseAI/Sources/Chat.swift b/FirebaseAI/Sources/Chat.swift
index 1aa2c3490c7..01d0b4a0056 100644
--- a/FirebaseAI/Sources/Chat.swift
+++ b/FirebaseAI/Sources/Chat.swift
@@ -146,11 +146,13 @@ public final class Chat: Sendable {
     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
+        guard !part.isThought else {
+          continue
+        }
 
-        default:
+        if let textPart = part as? TextPart {
+          combinedText += textPart.text
+        } else {
           // Don't combine it, just add to the content. If there's any text pending, add that as
           // a part.
           if !combinedText.isEmpty {
diff --git a/FirebaseAI/Sources/GenerateContentResponse.swift b/FirebaseAI/Sources/GenerateContentResponse.swift
index cb212e5a616..413b00aa632 100644
--- a/FirebaseAI/Sources/GenerateContentResponse.swift
+++ b/FirebaseAI/Sources/GenerateContentResponse.swift
@@ -66,12 +66,10 @@ public struct GenerateContentResponse: Sendable {
       return nil
     }
     let textValues: [String] = candidate.content.parts.compactMap { part in
-      switch part {
-      case let textPart as TextPart:
-        return textPart.text
-      default:
+      guard let textPart = part as? TextPart, !part.isThought else {
         return nil
       }
+      return textPart.text
     }
     guard textValues.count > 0 else {
       AILog.error(
@@ -89,12 +87,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 +103,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.

From 3173ece9761e153abe92a9f44f5c0e961fc5c478 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Mon, 14 Jul 2025 16:41:37 -0400
Subject: [PATCH 07/20] Add `thoughtSignature` to `Part`

---
 FirebaseAI/Sources/Chat.swift                 |  4 +-
 FirebaseAI/Sources/ModelContent.swift         | 73 +++++++++++++------
 .../Sources/Types/Internal/InternalPart.swift |  1 +
 FirebaseAI/Sources/Types/Public/Part.swift    | 37 +++++++---
 4 files changed, 80 insertions(+), 35 deletions(-)

diff --git a/FirebaseAI/Sources/Chat.swift b/FirebaseAI/Sources/Chat.swift
index 01d0b4a0056..361a5363236 100644
--- a/FirebaseAI/Sources/Chat.swift
+++ b/FirebaseAI/Sources/Chat.swift
@@ -150,7 +150,9 @@ public final class Chat: Sendable {
           continue
         }
 
-        if let textPart = part as? TextPart {
+        if let textPart = part as? TextPart,
+           // Parts with thought signatures must not be concatenated together.
+           part.thoughtSignature == nil {
           combinedText += textPart.text
         } else {
           // Don't combine it, just add to the content. If there's any text pending, add that as
diff --git a/FirebaseAI/Sources/ModelContent.swift b/FirebaseAI/Sources/ModelContent.swift
index 7fdd7d428c0..18a27617ebf 100644
--- a/FirebaseAI/Sources/ModelContent.swift
+++ b/FirebaseAI/Sources/ModelContent.swift
@@ -45,9 +45,12 @@ struct InternalPart: Equatable, Sendable {
 
   let isThought: Bool?
 
-  init(_ data: OneOfData, isThought: Bool?) {
+  let thoughtSignature: String?
+
+  init(_ data: OneOfData, isThought: Bool?, thoughtSignature: String?) {
     self.data = data
     self.isThought = isThought
+    self.thoughtSignature = thoughtSignature
   }
 }
 
@@ -66,15 +69,27 @@ public struct ModelContent: Equatable, Sendable {
     for part in internalParts {
       switch part.data {
       case let .text(text):
-        convertedParts.append(TextPart(text, isThought: part.isThought))
+        convertedParts.append(
+          TextPart(text, isThought: part.isThought, thoughtSignature: part.thoughtSignature)
+        )
       case let .inlineData(inlineData):
-        convertedParts.append(InlineDataPart(inlineData, isThought: part.isThought))
+        convertedParts.append(InlineDataPart(
+          inlineData, isThought: part.isThought, thoughtSignature: part.thoughtSignature
+        ))
       case let .fileData(fileData):
-        convertedParts.append(FileDataPart(fileData, isThought: part.isThought))
+        convertedParts.append(FileDataPart(
+          fileData,
+          isThought: part.isThought,
+          thoughtSignature: part.thoughtSignature
+        ))
       case let .functionCall(functionCall):
-        convertedParts.append(FunctionCallPart(functionCall, isThought: part.isThought))
+        convertedParts.append(FunctionCallPart(
+          functionCall, isThought: part.isThought, thoughtSignature: part.thoughtSignature
+        ))
       case let .functionResponse(functionResponse):
-        convertedParts.append(FunctionResponsePart(functionResponse, isThought: part.isThought))
+        convertedParts.append(FunctionResponsePart(
+          functionResponse, isThought: part.isThought, thoughtSignature: part.thoughtSignature
+        ))
       }
     }
     return convertedParts
@@ -90,28 +105,35 @@ public struct ModelContent: Equatable, Sendable {
     for part in parts {
       switch part {
       case let textPart as TextPart:
-        convertedParts.append(InternalPart(.text(textPart.text), isThought: textPart._isThought))
+        convertedParts.append(InternalPart(
+          .text(textPart.text),
+          isThought: textPart._isThought,
+          thoughtSignature: textPart.thoughtSignature
+        ))
       case let inlineDataPart as InlineDataPart:
-        convertedParts.append(
-          InternalPart(.inlineData(inlineDataPart.inlineData), isThought: inlineDataPart._isThought)
-        )
+        convertedParts.append(InternalPart(
+          .inlineData(inlineDataPart.inlineData),
+          isThought: inlineDataPart._isThought,
+          thoughtSignature: inlineDataPart.thoughtSignature
+        ))
       case let fileDataPart as FileDataPart:
-        convertedParts.append(
-          InternalPart(.fileData(fileDataPart.fileData), isThought: fileDataPart._isThought)
-        )
+        convertedParts.append(InternalPart(
+          .fileData(fileDataPart.fileData),
+          isThought: fileDataPart._isThought,
+          thoughtSignature: fileDataPart.thoughtSignature
+        ))
       case let functionCallPart as FunctionCallPart:
-        convertedParts.append(
-          InternalPart(
-            .functionCall(functionCallPart.functionCall), isThought: functionCallPart._isThought
-          )
-        )
+        convertedParts.append(InternalPart(
+          .functionCall(functionCallPart.functionCall),
+          isThought: functionCallPart._isThought,
+          thoughtSignature: functionCallPart.thoughtSignature
+        ))
       case let functionResponsePart as FunctionResponsePart:
-        convertedParts.append(
-          InternalPart(
-            .functionResponse(functionResponsePart.functionResponse),
-            isThought: functionResponsePart._isThought
-          )
-        )
+        convertedParts.append(InternalPart(
+          .functionResponse(functionResponsePart.functionResponse),
+          isThought: functionResponsePart._isThought,
+          thoughtSignature: functionResponsePart.thoughtSignature
+        ))
       default:
         fatalError()
       }
@@ -147,18 +169,21 @@ extension ModelContent: 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)
   }
 }
 
diff --git a/FirebaseAI/Sources/Types/Internal/InternalPart.swift b/FirebaseAI/Sources/Types/Internal/InternalPart.swift
index 062ae8aa93d..bb62dd4c0b5 100644
--- a/FirebaseAI/Sources/Types/Internal/InternalPart.swift
+++ b/FirebaseAI/Sources/Types/Internal/InternalPart.swift
@@ -68,6 +68,7 @@ struct ErrorPart: Part, Error {
   let error: Error
 
   let isThought = false
+  let thoughtSignature: String? = nil
 
   init(_ error: Error) {
     self.error = error
diff --git a/FirebaseAI/Sources/Types/Public/Part.swift b/FirebaseAI/Sources/Types/Public/Part.swift
index 4fea490a7bf..ca31f62f13e 100644
--- a/FirebaseAI/Sources/Types/Public/Part.swift
+++ b/FirebaseAI/Sources/Types/Public/Part.swift
@@ -20,6 +20,7 @@ import Foundation
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 public protocol Part: PartsRepresentable, Codable, Sendable, Equatable {
   var isThought: Bool { get }
+  var thoughtSignature: String? { get }
 }
 
 /// A text part containing a string value.
@@ -30,15 +31,18 @@ public struct TextPart: Part {
 
   public var isThought: Bool { _isThought ?? false }
 
+  public let thoughtSignature: String?
+
   let _isThought: Bool?
 
   public init(_ text: String) {
-    self.init(text, isThought: nil)
+    self.init(text, isThought: nil, thoughtSignature: nil)
   }
 
-  init(_ text: String, isThought: Bool?) {
+  init(_ text: String, isThought: Bool?, thoughtSignature: String?) {
     self.text = text
     _isThought = isThought
+    self.thoughtSignature = thoughtSignature
   }
 }
 
@@ -66,6 +70,8 @@ public struct InlineDataPart: Part {
 
   public var isThought: Bool { _isThought ?? false }
 
+  public let thoughtSignature: String?
+
   /// Creates an inline data part from data and a MIME type.
   ///
   /// > Important: Supported input types depend on the model on the model being used; see [input
@@ -81,12 +87,13 @@ public struct InlineDataPart: Part {
   ///     requirements](https://firebase.google.com/docs/vertex-ai/input-file-requirements) for
   ///     supported values.
   public init(data: Data, mimeType: String) {
-    self.init(InlineData(data: data, mimeType: mimeType), isThought: nil)
+    self.init(InlineData(data: data, mimeType: mimeType), isThought: nil, thoughtSignature: nil)
   }
 
-  init(_ inlineData: InlineData, isThought: Bool?) {
+  init(_ inlineData: InlineData, isThought: Bool?, thoughtSignature: String?) {
     self.inlineData = inlineData
     _isThought = isThought
+    self.thoughtSignature = thoughtSignature
   }
 }
 
@@ -99,6 +106,7 @@ public struct FileDataPart: Part {
   public var uri: String { fileData.fileURI }
   public var mimeType: String { fileData.mimeType }
   public var isThought: Bool { _isThought ?? false }
+  public let thoughtSignature: String?
 
   /// Constructs a new file data part.
   ///
@@ -110,12 +118,13 @@ public struct FileDataPart: Part {
   ///     requirements](https://firebase.google.com/docs/vertex-ai/input-file-requirements) for
   ///     supported values.
   public init(uri: String, mimeType: String) {
-    self.init(FileData(fileURI: uri, mimeType: mimeType), isThought: nil)
+    self.init(FileData(fileURI: uri, mimeType: mimeType), isThought: nil, thoughtSignature: nil)
   }
 
-  init(_ fileData: FileData, isThought: Bool?) {
+  init(_ fileData: FileData, isThought: Bool?, thoughtSignature: String?) {
     self.fileData = fileData
     _isThought = isThought
+    self.thoughtSignature = thoughtSignature
   }
 }
 
@@ -133,6 +142,8 @@ public struct FunctionCallPart: Part {
 
   public var isThought: Bool { _isThought ?? false }
 
+  public let thoughtSignature: String?
+
   /// Constructs a new function call part.
   ///
   /// > Note: A `FunctionCallPart` is typically received from the model, rather than created
@@ -142,12 +153,13 @@ public struct FunctionCallPart: Part {
   ///   - name: The name of the function to call.
   ///   - args: The function parameters and values.
   public init(name: String, args: JSONObject) {
-    self.init(FunctionCall(name: name, args: args), isThought: nil)
+    self.init(FunctionCall(name: name, args: args), isThought: nil, thoughtSignature: nil)
   }
 
-  init(_ functionCall: FunctionCall, isThought: Bool?) {
+  init(_ functionCall: FunctionCall, isThought: Bool?, thoughtSignature: String?) {
     self.functionCall = functionCall
     _isThought = isThought
+    self.thoughtSignature = thoughtSignature
   }
 }
 
@@ -169,17 +181,22 @@ public struct FunctionResponsePart: Part {
 
   public var isThought: Bool { _isThought ?? false }
 
+  public let thoughtSignature: String?
+
   /// Constructs a new `FunctionResponse`.
   ///
   /// - Parameters:
   ///   - name: The name of the function that was called.
   ///   - response: The function's response.
   public init(name: String, response: JSONObject) {
-    self.init(FunctionResponse(name: name, response: response), isThought: nil)
+    self.init(
+      FunctionResponse(name: name, response: response), isThought: nil, thoughtSignature: nil
+    )
   }
 
-  init(_ functionResponse: FunctionResponse, isThought: Bool?) {
+  init(_ functionResponse: FunctionResponse, isThought: Bool?, thoughtSignature: String?) {
     self.functionResponse = functionResponse
     _isThought = isThought
+    self.thoughtSignature = thoughtSignature
   }
 }

From 16eb6497b27f44f3bc49bf9689c3e2a276a3ff40 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Tue, 15 Jul 2025 18:11:10 -0400
Subject: [PATCH 08/20] Handle part aggregation in `Chat` with "thoughts"

---
 FirebaseAI/Sources/Chat.swift | 30 +++++++++++++++++++++++-------
 1 file changed, 23 insertions(+), 7 deletions(-)

diff --git a/FirebaseAI/Sources/Chat.swift b/FirebaseAI/Sources/Chat.swift
index 361a5363236..3caede2e428 100644
--- a/FirebaseAI/Sources/Chat.swift
+++ b/FirebaseAI/Sources/Chat.swift
@@ -142,31 +142,47 @@ public final class Chat: Sendable {
 
   private func aggregatedChunks(_ chunks: [ModelContent]) -> ModelContent {
     var parts: [any Part] = []
+    var previousPart = chunks.first?.parts.first
     var combinedText = ""
+    var combinedThoughts = ""
     for aggregate in chunks {
       // Loop through all the parts, aggregating the text and adding the images.
       for part in aggregate.parts {
-        guard !part.isThought else {
-          continue
-        }
-
         if let textPart = part as? TextPart,
+           // Thought summaries must not be combined with regular output.
+           (previousPart?.isThought ?? false) == part.isThought,
            // Parts with thought signatures must not be concatenated together.
            part.thoughtSignature == nil {
-          combinedText += textPart.text
+          if part.isThought {
+            combinedThoughts += textPart.text
+          } else {
+            combinedText += textPart.text
+          }
         } else {
-          // Don't combine it, just add to the content. If there's any text pending, add that as
-          // a part.
+          // This is a non-text part (e.g., inline data or a function call).
+
+          // 1. Append any thought summaries we've accumulated so far.
+          if !combinedThoughts.isEmpty {
+            parts.append(TextPart(combinedThoughts))
+            combinedThoughts = ""
+          }
+          // 2. Append any text we've accumulated so far.
           if !combinedText.isEmpty {
             parts.append(TextPart(combinedText))
             combinedText = ""
           }
 
+          // Then append the non-text part itself.
           parts.append(part)
         }
+
+        previousPart = part
       }
     }
 
+    if !combinedThoughts.isEmpty {
+      parts.append(TextPart(combinedThoughts))
+    }
     if !combinedText.isEmpty {
       parts.append(TextPart(combinedText))
     }

From bddbee60d7ff2b20cab179cefffdfe7bf21ae37e Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Tue, 15 Jul 2025 18:22:05 -0400
Subject: [PATCH 09/20] Fix bug and simplify `aggregatedChunks` with nested
 `flush` function

---
 FirebaseAI/Sources/Chat.swift | 65 +++++++++++++++++------------------
 1 file changed, 31 insertions(+), 34 deletions(-)

diff --git a/FirebaseAI/Sources/Chat.swift b/FirebaseAI/Sources/Chat.swift
index 3caede2e428..6df94ced839 100644
--- a/FirebaseAI/Sources/Chat.swift
+++ b/FirebaseAI/Sources/Chat.swift
@@ -142,50 +142,47 @@ public final class Chat: Sendable {
 
   private func aggregatedChunks(_ chunks: [ModelContent]) -> ModelContent {
     var parts: [any Part] = []
-    var previousPart = chunks.first?.parts.first
     var combinedText = ""
     var combinedThoughts = ""
-    for aggregate in chunks {
-      // Loop through all the parts, aggregating the text and adding the images.
-      for part in aggregate.parts {
-        if let textPart = part as? TextPart,
-           // Thought summaries must not be combined with regular output.
-           (previousPart?.isThought ?? false) == part.isThought,
-           // Parts with thought signatures must not be concatenated together.
-           part.thoughtSignature == nil {
-          if part.isThought {
-            combinedThoughts += textPart.text
-          } else {
-            combinedText += textPart.text
+
+    func flush() {
+      if !combinedThoughts.isEmpty {
+        parts.append(TextPart(combinedThoughts))
+        combinedThoughts = ""
+      }
+      if !combinedText.isEmpty {
+        parts.append(TextPart(combinedText))
+        combinedText = ""
+      }
+    }
+
+    // Loop through all the parts, aggregating the text.
+    for part in chunks.flatMap({ $0.parts }) {
+      // Only text parts may be combined.
+      if let textPart = part as? TextPart, part.thoughtSignature == nil {
+        // Thought summaries must not be combined with regular text.
+        if textPart.isThought {
+          // If we were combining regular text, flush it before handling "thoughts".
+          if !combinedText.isEmpty {
+            flush()
           }
+          combinedThoughts += textPart.text
         } else {
-          // This is a non-text part (e.g., inline data or a function call).
-
-          // 1. Append any thought summaries we've accumulated so far.
+          // If we were combining "thoughts", flush it before handling regular text.
           if !combinedThoughts.isEmpty {
-            parts.append(TextPart(combinedThoughts))
-            combinedThoughts = ""
+            flush()
           }
-          // 2. Append any text we've accumulated so far.
-          if !combinedText.isEmpty {
-            parts.append(TextPart(combinedText))
-            combinedText = ""
-          }
-
-          // Then append the non-text part itself.
-          parts.append(part)
+          combinedText += textPart.text
         }
-
-        previousPart = part
+      } else {
+        // This is a non-combinable part (not text), flush any pending text.
+        flush()
+        parts.append(part)
       }
     }
 
-    if !combinedThoughts.isEmpty {
-      parts.append(TextPart(combinedThoughts))
-    }
-    if !combinedText.isEmpty {
-      parts.append(TextPart(combinedText))
-    }
+    // Flush any remaining text.
+    flush()
 
     return ModelContent(role: "model", parts: parts)
   }

From 44100f9c5df7974db977359d634a63537ec182d6 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Fri, 18 Jul 2025 15:56:43 -0400
Subject: [PATCH 10/20] Refactor `thoughtSignature` to be `internal`

---
 FirebaseAI/Sources/Chat.swift              | 16 ++++++++--------
 FirebaseAI/Sources/ModelContent.swift      |  5 +++++
 FirebaseAI/Sources/Types/Public/Part.swift | 13 +++++--------
 3 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/FirebaseAI/Sources/Chat.swift b/FirebaseAI/Sources/Chat.swift
index bdf4ee8eb1c..80e908a8f57 100644
--- a/FirebaseAI/Sources/Chat.swift
+++ b/FirebaseAI/Sources/Chat.swift
@@ -147,38 +147,38 @@ public final class Chat: Sendable {
   }
 
   private func aggregatedChunks(_ chunks: [ModelContent]) -> ModelContent {
-    var parts: [any Part] = []
+    var parts: [InternalPart] = []
     var combinedText = ""
     var combinedThoughts = ""
 
     func flush() {
       if !combinedThoughts.isEmpty {
-        parts.append(TextPart(combinedThoughts))
+        parts.append(InternalPart(.text(combinedThoughts), isThought: true, thoughtSignature: nil))
         combinedThoughts = ""
       }
       if !combinedText.isEmpty {
-        parts.append(TextPart(combinedText))
+        parts.append(InternalPart(.text(combinedText), isThought: nil, thoughtSignature: nil))
         combinedText = ""
       }
     }
 
     // Loop through all the parts, aggregating the text.
-    for part in chunks.flatMap({ $0.parts }) {
+    for part in chunks.flatMap({ $0.internalParts }) {
       // Only text parts may be combined.
-      if let textPart = part as? TextPart, part.thoughtSignature == nil {
+      if case let .text(text) = part.data, part.thoughtSignature == nil {
         // Thought summaries must not be combined with regular text.
-        if textPart.isThought {
+        if part.isThought ?? false {
           // If we were combining regular text, flush it before handling "thoughts".
           if !combinedText.isEmpty {
             flush()
           }
-          combinedThoughts += textPart.text
+          combinedThoughts += text
         } else {
           // If we were combining "thoughts", flush it before handling regular text.
           if !combinedThoughts.isEmpty {
             flush()
           }
-          combinedText += textPart.text
+          combinedText += text
         }
       } else {
         // This is a non-combinable part (not text), flush any pending text.
diff --git a/FirebaseAI/Sources/ModelContent.swift b/FirebaseAI/Sources/ModelContent.swift
index 18a27617ebf..558c30b2789 100644
--- a/FirebaseAI/Sources/ModelContent.swift
+++ b/FirebaseAI/Sources/ModelContent.swift
@@ -147,6 +147,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
diff --git a/FirebaseAI/Sources/Types/Public/Part.swift b/FirebaseAI/Sources/Types/Public/Part.swift
index ca31f62f13e..2c9006fb1ab 100644
--- a/FirebaseAI/Sources/Types/Public/Part.swift
+++ b/FirebaseAI/Sources/Types/Public/Part.swift
@@ -20,7 +20,6 @@ import Foundation
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 public protocol Part: PartsRepresentable, Codable, Sendable, Equatable {
   var isThought: Bool { get }
-  var thoughtSignature: String? { get }
 }
 
 /// A text part containing a string value.
@@ -31,7 +30,7 @@ public struct TextPart: Part {
 
   public var isThought: Bool { _isThought ?? false }
 
-  public let thoughtSignature: String?
+  let thoughtSignature: String?
 
   let _isThought: Bool?
 
@@ -70,7 +69,7 @@ public struct InlineDataPart: Part {
 
   public var isThought: Bool { _isThought ?? false }
 
-  public let thoughtSignature: String?
+  let thoughtSignature: String?
 
   /// Creates an inline data part from data and a MIME type.
   ///
@@ -102,11 +101,11 @@ public struct InlineDataPart: Part {
 public struct FileDataPart: Part {
   let fileData: FileData
   let _isThought: Bool?
+  let thoughtSignature: String?
 
   public var uri: String { fileData.fileURI }
   public var mimeType: String { fileData.mimeType }
   public var isThought: Bool { _isThought ?? false }
-  public let thoughtSignature: String?
 
   /// Constructs a new file data part.
   ///
@@ -133,6 +132,7 @@ public struct FileDataPart: Part {
 public struct FunctionCallPart: Part {
   let functionCall: FunctionCall
   let _isThought: Bool?
+  let thoughtSignature: String?
 
   /// The name of the function to call.
   public var name: String { functionCall.name }
@@ -142,8 +142,6 @@ public struct FunctionCallPart: Part {
 
   public var isThought: Bool { _isThought ?? false }
 
-  public let thoughtSignature: String?
-
   /// Constructs a new function call part.
   ///
   /// > Note: A `FunctionCallPart` is typically received from the model, rather than created
@@ -172,6 +170,7 @@ public struct FunctionCallPart: Part {
 public struct FunctionResponsePart: Part {
   let functionResponse: FunctionResponse
   let _isThought: Bool?
+  let thoughtSignature: String?
 
   /// The name of the function that was called.
   public var name: String { functionResponse.name }
@@ -181,8 +180,6 @@ public struct FunctionResponsePart: Part {
 
   public var isThought: Bool { _isThought ?? false }
 
-  public let thoughtSignature: String?
-
   /// Constructs a new `FunctionResponse`.
   ///
   /// - Parameters:

From 835748fb874664184d49c1fb7900c742ac144ecb Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Fri, 18 Jul 2025 18:02:57 -0400
Subject: [PATCH 11/20] Add `thoughtSummary` convenience accessor

---
 .../Sources/GenerateContentResponse.swift     | 50 +++++++++++--------
 1 file changed, 29 insertions(+), 21 deletions(-)

diff --git a/FirebaseAI/Sources/GenerateContentResponse.swift b/FirebaseAI/Sources/GenerateContentResponse.swift
index 19e4d797419..5b3eb614e32 100644
--- a/FirebaseAI/Sources/GenerateContentResponse.swift
+++ b/FirebaseAI/Sources/GenerateContentResponse.swift
@@ -58,27 +58,11 @@ public struct GenerateContentResponse: Sendable {
 
   /// The response's content as text, if it exists.
   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
-      guard let textPart = part as? TextPart, !part.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: " ")
+    return text(isThought: false)
+  }
+
+  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.
@@ -118,6 +102,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

From fcdee03ceb72eda65923cee509ababa629cdeb4b Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Mon, 28 Jul 2025 16:06:51 -0400
Subject: [PATCH 12/20] Add docs for `includeThoughts` and `isThought`

---
 FirebaseAI/Sources/Types/Public/Part.swift           | 5 +++++
 FirebaseAI/Sources/Types/Public/ThinkingConfig.swift | 9 +++++++++
 2 files changed, 14 insertions(+)

diff --git a/FirebaseAI/Sources/Types/Public/Part.swift b/FirebaseAI/Sources/Types/Public/Part.swift
index 2c9006fb1ab..fb743d1025d 100644
--- a/FirebaseAI/Sources/Types/Public/Part.swift
+++ b/FirebaseAI/Sources/Types/Public/Part.swift
@@ -19,6 +19,11 @@ import Foundation
 /// Within a single value of ``Part``, different data types may not mix.
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 public protocol Part: PartsRepresentable, Codable, Sendable, Equatable {
+  /// Indicates whether this `Part` is a summary of the model's internal thinking process.
+  ///
+  /// When `includeThoughts` is set to `true` in ``ThinkingConfig``, the model may return one or
+  /// more "thought" parts that provide insight into how it reasoned through the prompt to arrive
+  /// at the final answer. These parts will have `isThought` set to `true`.
   var isThought: Bool { get }
 }
 
diff --git a/FirebaseAI/Sources/Types/Public/ThinkingConfig.swift b/FirebaseAI/Sources/Types/Public/ThinkingConfig.swift
index c9e2cc43f4e..a339f8fa1d1 100644
--- a/FirebaseAI/Sources/Types/Public/ThinkingConfig.swift
+++ b/FirebaseAI/Sources/Types/Public/ThinkingConfig.swift
@@ -37,12 +37,21 @@ public struct ThinkingConfig: Sendable {
   /// feature or if the specified budget is not within the model's supported range.
   let thinkingBudget: Int?
 
+  /// Whether summaries of the model's "thoughts" are included in responses.
+  ///
+  /// When `includeThoughts` is set to `true`, the model will return a summary of its internal
+  /// thinking process alongside the final answer. This can provide valuable insight into how the
+  /// model arrived at its conclusion, which is particularly useful for complex or creative tasks.
+  ///
+  /// If you don't specify a value for `includeThoughts` (`nil`), the model will use its default
+  /// behavior (which is typically to not include thought summaries).
   let includeThoughts: Bool?
 
   /// Initializes a new `ThinkingConfig`.
   ///
   /// - Parameters:
   ///   - thinkingBudget: The maximum number of tokens to be used for the model's thinking process.
+  ///   - includeThoughts: If true, summaries of the model's "thoughts" are included in responses.
   public init(thinkingBudget: Int? = nil, includeThoughts: Bool? = nil) {
     self.thinkingBudget = thinkingBudget
     self.includeThoughts = includeThoughts

From 98d14cbe0eb8e135035ddf2e77000997d63b43b0 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Mon, 28 Jul 2025 22:16:47 -0400
Subject: [PATCH 13/20] Add mock response tests to
 `GenerativeModelVertexAITests`

---
 .../Unit/GenerativeModelVertexAITests.swift   | 50 +++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift b/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift
index 6557735ccc4..7c23726f152 100644
--- a/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift
+++ b/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift
@@ -434,6 +434,29 @@ final class GenerativeModelVertexAITests: XCTestCase {
     XCTAssertEqual(text, "The sum of [1, 2, 3] is")
   }
 
+  func testGenerateContent_success_thinking_thoughtSummary() async throws {
+    MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+      forResource: "unary-success-thinking-reply-thought-summary",
+      withExtension: "json",
+      subdirectory: vertexSubdirectory
+    )
+
+    let response = try await model.generateContent(testPrompt)
+
+    XCTAssertEqual(response.candidates.count, 1)
+    let candidate = try XCTUnwrap(response.candidates.first)
+    XCTAssertEqual(candidate.finishReason, .stop)
+    XCTAssertEqual(candidate.content.parts.count, 2)
+    let thoughtPart = try XCTUnwrap(candidate.content.parts.first as? TextPart)
+    XCTAssertTrue(thoughtPart.isThought)
+    XCTAssertTrue(thoughtPart.text.hasPrefix("Right, someone needs the city where Google"))
+    XCTAssertEqual(response.thoughtSummary, thoughtPart.text)
+    let textPart = try XCTUnwrap(candidate.content.parts.last as? TextPart)
+    XCTAssertFalse(textPart.isThought)
+    XCTAssertEqual(textPart.text, "Mountain View")
+    XCTAssertEqual(response.text, textPart.text)
+  }
+
   func testGenerateContent_success_image_invalidSafetyRatingsIgnored() async throws {
     MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
       forResource: "unary-success-image-invalid-safety-ratings",
@@ -1330,6 +1353,33 @@ final class GenerativeModelVertexAITests: XCTestCase {
     XCTAssertFalse(citations.contains { $0.license?.isEmpty ?? false })
   }
 
+  func testGenerateContentStream_successWithThinking_thoughtSummary() async throws {
+    MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+      forResource: "streaming-success-thinking-reply-thought-summary",
+      withExtension: "txt",
+      subdirectory: vertexSubdirectory
+    )
+
+    var thoughtSummary = ""
+    var text = ""
+    let stream = try model.generateContentStream("Hi")
+    for try await response in stream {
+      let candidate = try XCTUnwrap(response.candidates.first)
+      XCTAssertEqual(candidate.content.parts.count, 1)
+      let part = try XCTUnwrap(candidate.content.parts.first)
+      let textPart = try XCTUnwrap(part as? TextPart)
+      if textPart.isThought {
+        let newThought = try XCTUnwrap(response.thoughtSummary)
+        thoughtSummary.append(newThought)
+      } else {
+        text.append(textPart.text)
+      }
+    }
+
+    XCTAssertTrue(thoughtSummary.hasPrefix("**Understanding the Core Question**"))
+    XCTAssertTrue(text.hasPrefix("The sky is blue due to a phenomenon"))
+  }
+
   func testGenerateContentStream_successWithInvalidSafetyRatingsIgnored() async throws {
     MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
       forResource: "streaming-success-image-invalid-safety-ratings",

From 0a7329f8c1f73b31c63be2bc581a5392cff8a855 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Mon, 28 Jul 2025 22:59:00 -0400
Subject: [PATCH 14/20] Add mock response tests to
 `GenerativeModelGoogleAITests`

---
 .../Unit/GenerativeModelGoogleAITests.swift   | 112 ++++++++++++++++++
 1 file changed, 112 insertions(+)

diff --git a/FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift b/FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift
index 103943e6f92..00e0d398855 100644
--- a/FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift
+++ b/FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift
@@ -262,6 +262,52 @@ final class GenerativeModelGoogleAITests: XCTestCase {
     )
   }
 
+  func testGenerateContent_success_thinking_thoughtSummary() async throws {
+    MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+      forResource: "unary-success-thinking-reply-thought-summary",
+      withExtension: "json",
+      subdirectory: googleAISubdirectory
+    )
+
+    let response = try await model.generateContent(testPrompt)
+
+    XCTAssertEqual(response.candidates.count, 1)
+    let candidate = try XCTUnwrap(response.candidates.first)
+    XCTAssertEqual(candidate.content.parts.count, 2)
+    let thoughtPart = try XCTUnwrap(candidate.content.parts.first as? TextPart)
+    XCTAssertTrue(thoughtPart.isThought)
+    XCTAssertTrue(thoughtPart.text.hasPrefix("**Thinking About Google's Headquarters**"))
+    XCTAssertEqual(thoughtPart.text, response.thoughtSummary)
+    let textPart = try XCTUnwrap(candidate.content.parts.last as? TextPart)
+    XCTAssertFalse(textPart.isThought)
+    XCTAssertEqual(textPart.text, "Mountain View")
+    XCTAssertEqual(textPart.text, response.text)
+  }
+
+  func testGenerateContent_success_thinking_functionCall_thoughtSummaryAndSignature() async throws {
+    MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+      forResource: "unary-success-thinking-function-call-thought-summary-signature",
+      withExtension: "json",
+      subdirectory: googleAISubdirectory
+    )
+
+    let response = try await model.generateContent(testPrompt)
+
+    XCTAssertEqual(response.candidates.count, 1)
+    let candidate = try XCTUnwrap(response.candidates.first)
+    XCTAssertEqual(candidate.finishReason, .stop)
+    XCTAssertEqual(candidate.content.parts.count, 2)
+    let thoughtPart = try XCTUnwrap(candidate.content.parts.first as? TextPart)
+    XCTAssertTrue(thoughtPart.isThought)
+    XCTAssertTrue(thoughtPart.text.hasPrefix("**Thinking Through the New Year's Eve Calculation**"))
+    let functionCallPart = try XCTUnwrap(candidate.content.parts.last as? FunctionCallPart)
+    XCTAssertFalse(functionCallPart.isThought)
+    XCTAssertEqual(functionCallPart.name, "now")
+    XCTAssertTrue(functionCallPart.args.isEmpty)
+    let thoughtSignature = try XCTUnwrap(functionCallPart.thoughtSignature)
+    XCTAssertTrue(thoughtSignature.hasPrefix("CtQOAVSoXO74PmYr9AFu"))
+  }
+
   func testGenerateContent_failure_invalidAPIKey() async throws {
     let expectedStatusCode = 400
     MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
@@ -397,6 +443,72 @@ final class GenerativeModelGoogleAITests: XCTestCase {
     XCTAssertNil(citation.publicationDate)
   }
 
+  func testGenerateContentStream_successWithThoughtSummary() async throws {
+    MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+      forResource: "streaming-success-thinking-reply-thought-summary",
+      withExtension: "txt",
+      subdirectory: googleAISubdirectory
+    )
+
+    var thoughtSummary = ""
+    var text = ""
+    let stream = try model.generateContentStream("Hi")
+    for try await response in stream {
+      let candidate = try XCTUnwrap(response.candidates.first)
+      XCTAssertEqual(candidate.content.parts.count, 1)
+      let textPart = try XCTUnwrap(candidate.content.parts.first as? TextPart)
+      if textPart.isThought {
+        let newThought = try XCTUnwrap(response.thoughtSummary)
+        XCTAssertEqual(textPart.text, newThought)
+        thoughtSummary.append(newThought)
+      } else {
+        let newText = try XCTUnwrap(response.text)
+        XCTAssertEqual(textPart.text, newText)
+        text.append(newText)
+      }
+    }
+
+    XCTAssertTrue(thoughtSummary.hasPrefix("**Exploring Sky Color**"))
+    XCTAssertTrue(text.hasPrefix("The sky is blue because"))
+  }
+
+  func testGenerateContentStream_success_thinking_functionCall_thoughtSummary_signature() async throws {
+    MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+      forResource: "streaming-success-thinking-function-call-thought-summary-signature",
+      withExtension: "txt",
+      subdirectory: googleAISubdirectory
+    )
+
+    var thoughtSummary = ""
+    var functionCalls: [FunctionCallPart] = []
+    let stream = try model.generateContentStream("Hi")
+    for try await response in stream {
+      let candidate = try XCTUnwrap(response.candidates.first)
+      XCTAssertEqual(candidate.content.parts.count, 1)
+      let part = try XCTUnwrap(candidate.content.parts.first)
+      if part.isThought {
+        let textPart = try XCTUnwrap(part as? TextPart)
+        let newThought = try XCTUnwrap(response.thoughtSummary)
+        XCTAssertEqual(textPart.text, newThought)
+        thoughtSummary.append(newThought)
+      } else {
+        let functionCallPart = try XCTUnwrap(part as? FunctionCallPart)
+        XCTAssertEqual(response.functionCalls.count, 1)
+        let newFunctionCall = try XCTUnwrap(response.functionCalls.first)
+        XCTAssertEqual(functionCallPart, newFunctionCall)
+        functionCalls.append(newFunctionCall)
+      }
+    }
+
+    XCTAssertTrue(thoughtSummary.hasPrefix("**Calculating the Days**"))
+    XCTAssertEqual(functionCalls.count, 1)
+    let functionCall = try XCTUnwrap(functionCalls.first)
+    XCTAssertEqual(functionCall.name, "now")
+    XCTAssertTrue(functionCall.args.isEmpty)
+    let thoughtSignature = try XCTUnwrap(functionCall.thoughtSignature)
+    XCTAssertTrue(thoughtSignature.hasPrefix("CiIBVKhc7vB+vaaq6rA"))
+  }
+
   func testGenerateContentStream_failureInvalidAPIKey() async throws {
     MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
       forResource: "unary-failure-api-key",

From 959e2b8b45b3ab70eabdbcb26b759bb07312f5cf Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Tue, 29 Jul 2025 15:39:31 -0400
Subject: [PATCH 15/20] Add `generateContentThinkingFunctionCalling`
 integration test

---
 .../GenerateContentIntegrationTests.swift     | 122 ++++++++++++++++++
 1 file changed, 122 insertions(+)

diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
index 0092efe33ce..d662e9c61a3 100644
--- a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
+++ b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
@@ -231,6 +231,128 @@ struct GenerateContentIntegrationTests {
     ))
   }
 
+  @Test(
+    arguments: [
+      (.vertexAI_v1beta_global, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: 0)),
+      (.vertexAI_v1beta_global, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: 24576)),
+      (.vertexAI_v1beta_global, ModelNames.gemini2_5_Flash, ThinkingConfig(
+        thinkingBudget: 24576, includeThoughts: true
+      )),
+      (.vertexAI_v1beta_global, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: 128)),
+      (.vertexAI_v1beta_global, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: 32768)),
+      (.vertexAI_v1beta_global, ModelNames.gemini2_5_Pro, ThinkingConfig(
+        thinkingBudget: 32768, includeThoughts: true
+      )),
+      (.googleAI_v1beta, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: 0)),
+      (.googleAI_v1beta, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: 24576)),
+      (.googleAI_v1beta, ModelNames.gemini2_5_Flash, ThinkingConfig(
+        thinkingBudget: 24576, includeThoughts: true
+      )),
+      (.googleAI_v1beta, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: 128)),
+      (.googleAI_v1beta, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: 32768)),
+      (.googleAI_v1beta, ModelNames.gemini2_5_Pro, ThinkingConfig(
+        thinkingBudget: 32768, includeThoughts: true
+      )),
+    ] as [(InstanceConfig, String, ThinkingConfig)]
+  )
+  func generateContentThinkingFunctionCalling(_ config: InstanceConfig, modelName: String,
+                                              thinkingConfig: ThinkingConfig) async throws {
+    let currentLocationDeclaration = FunctionDeclaration(
+      name: "currentLocation",
+      description: "Returns the user's current city, province or state, and country",
+      parameters: [:]
+    )
+    let getTemperatureDeclaration = FunctionDeclaration(
+      name: "getTemperature",
+      description: "Returns the current temperature in Celsius for the specified location",
+      parameters: [
+        "city": .string(),
+        "region": .string(description: "The province or state"),
+        "country": .string(),
+      ]
+    )
+    let model = FirebaseAI.componentInstance(config).generativeModel(
+      modelName: modelName,
+      generationConfig: GenerationConfig(
+        temperature: 0.0,
+        topP: 0.0,
+        topK: 1,
+        thinkingConfig: thinkingConfig
+      ),
+      safetySettings: safetySettings,
+      tools: [.functionDeclarations([currentLocationDeclaration, getTemperatureDeclaration])]
+    )
+    let chat = model.startChat()
+    let prompt = """
+    What is the temperature outside right now? Respond in the format:
+    - Location: City, Province/State, Country
+    - Temperature: #C
+
+    Example Output:
+    - Location: Vancouver, British Columbia, Canada
+    - Temperature: 15C
+    """
+
+    let response = try await chat.sendMessage(prompt)
+
+    var thoughtSignatureCount = 0
+    #expect(response.functionCalls.count == 1)
+    let locationFunctionCall = try #require(response.functionCalls.first)
+    try #require(locationFunctionCall.name == currentLocationDeclaration.name)
+    #expect(locationFunctionCall.args.isEmpty)
+    #expect(locationFunctionCall.isThought == false)
+    if locationFunctionCall.thoughtSignature != nil {
+      thoughtSignatureCount += 1
+    }
+
+    let locationFunctionResponse = FunctionResponsePart(
+      name: locationFunctionCall.name,
+      response: [
+        "city": .string("Waterloo"),
+        "province": .string("Ontario"),
+        "country": .string("Canada"),
+      ]
+    )
+
+    let response2 = try await chat.sendMessage(locationFunctionResponse)
+
+    #expect(response2.functionCalls.count == 1)
+    let temperatureFunctionCall = try #require(response2.functionCalls.first)
+    try #require(temperatureFunctionCall.name == getTemperatureDeclaration.name)
+    #expect(temperatureFunctionCall.args == [
+      "city": .string("Waterloo"),
+      "region": .string("Ontario"),
+      "country": .string("Canada"),
+    ])
+    #expect(temperatureFunctionCall.isThought == false)
+    if temperatureFunctionCall.thoughtSignature != nil {
+      thoughtSignatureCount += 1
+    }
+
+    let temperatureFunctionResponse = FunctionResponsePart(
+      name: locationFunctionCall.name,
+      response: [
+        "temperature": .number(25),
+        "units": .string("Celsius"),
+      ]
+    )
+
+    let response3 = try await chat.sendMessage(temperatureFunctionResponse)
+
+    #expect(response3.functionCalls.isEmpty)
+    let finalText = try #require(response3.text).trimmingCharacters(in: .whitespacesAndNewlines)
+    #expect(finalText == """
+    - Location: Waterloo, Ontario, Canada
+    - Temperature: 25C
+    """)
+
+    if let _ = thinkingConfig.includeThoughts, case .googleAI = config.apiConfig.service {
+      #expect(thoughtSignatureCount > 0)
+    } else {
+      #expect(thoughtSignatureCount == 0)
+    }
+  }
+
   @Test(arguments: [
     InstanceConfig.vertexAI_v1beta,
     InstanceConfig.vertexAI_v1beta_global,

From b142c34fddfdc94beb3ba1a3a79166fad68a2962 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Tue, 29 Jul 2025 16:04:57 -0400
Subject: [PATCH 16/20] Use dynamic thinking budget in
 `generateContentThinkingFunctionCalling`

---
 .../GenerateContentIntegrationTests.swift     | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
index d662e9c61a3..ac2a47333bf 100644
--- a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
+++ b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
@@ -234,24 +234,23 @@ struct GenerateContentIntegrationTests {
   @Test(
     arguments: [
       (.vertexAI_v1beta_global, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: 0)),
-      (.vertexAI_v1beta_global, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: 24576)),
+      (.vertexAI_v1beta_global, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: -1)),
       (.vertexAI_v1beta_global, ModelNames.gemini2_5_Flash, ThinkingConfig(
-        thinkingBudget: 24576, includeThoughts: true
+        thinkingBudget: -1, includeThoughts: true
       )),
-      (.vertexAI_v1beta_global, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: 128)),
-      (.vertexAI_v1beta_global, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: 32768)),
+      (.vertexAI_v1beta_global, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: -1)),
       (.vertexAI_v1beta_global, ModelNames.gemini2_5_Pro, ThinkingConfig(
-        thinkingBudget: 32768, includeThoughts: true
+        thinkingBudget: -1, includeThoughts: true
       )),
       (.googleAI_v1beta, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: 0)),
-      (.googleAI_v1beta, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: 24576)),
+      (.googleAI_v1beta, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: -1)),
       (.googleAI_v1beta, ModelNames.gemini2_5_Flash, ThinkingConfig(
-        thinkingBudget: 24576, includeThoughts: true
+        thinkingBudget: -1, includeThoughts: true
       )),
-      (.googleAI_v1beta, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: 128)),
-      (.googleAI_v1beta, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: 32768)),
+      (.googleAI_v1beta, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: -1)),
+      (.googleAI_v1beta, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: -1)),
       (.googleAI_v1beta, ModelNames.gemini2_5_Pro, ThinkingConfig(
-        thinkingBudget: 32768, includeThoughts: true
+        thinkingBudget: -1, includeThoughts: true
       )),
     ] as [(InstanceConfig, String, ThinkingConfig)]
   )

From 86a725bc217d4a23adc6c61832af07b60910ed25 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Wed, 30 Jul 2025 18:55:42 -0400
Subject: [PATCH 17/20] Try to reduce flakiness of
 `generateContentThinkingFunctionCalling` test

---
 .../GenerateContentIntegrationTests.swift     | 36 +++++++++----------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
index ac2a47333bf..b672356a84d 100644
--- a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
+++ b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
@@ -233,7 +233,6 @@ struct GenerateContentIntegrationTests {
 
   @Test(
     arguments: [
-      (.vertexAI_v1beta_global, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: 0)),
       (.vertexAI_v1beta_global, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: -1)),
       (.vertexAI_v1beta_global, ModelNames.gemini2_5_Flash, ThinkingConfig(
         thinkingBudget: -1, includeThoughts: true
@@ -242,13 +241,11 @@ struct GenerateContentIntegrationTests {
       (.vertexAI_v1beta_global, ModelNames.gemini2_5_Pro, ThinkingConfig(
         thinkingBudget: -1, includeThoughts: true
       )),
-      (.googleAI_v1beta, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: 0)),
       (.googleAI_v1beta, ModelNames.gemini2_5_Flash, ThinkingConfig(thinkingBudget: -1)),
       (.googleAI_v1beta, ModelNames.gemini2_5_Flash, ThinkingConfig(
         thinkingBudget: -1, includeThoughts: true
       )),
       (.googleAI_v1beta, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: -1)),
-      (.googleAI_v1beta, ModelNames.gemini2_5_Pro, ThinkingConfig(thinkingBudget: -1)),
       (.googleAI_v1beta, ModelNames.gemini2_5_Pro, ThinkingConfig(
         thinkingBudget: -1, includeThoughts: true
       )),
@@ -257,7 +254,7 @@ struct GenerateContentIntegrationTests {
   func generateContentThinkingFunctionCalling(_ config: InstanceConfig, modelName: String,
                                               thinkingConfig: ThinkingConfig) async throws {
     let currentLocationDeclaration = FunctionDeclaration(
-      name: "currentLocation",
+      name: "getCurrentLocation",
       description: "Returns the user's current city, province or state, and country",
       parameters: [:]
     )
@@ -279,18 +276,23 @@ struct GenerateContentIntegrationTests {
         thinkingConfig: thinkingConfig
       ),
       safetySettings: safetySettings,
-      tools: [.functionDeclarations([currentLocationDeclaration, getTemperatureDeclaration])]
+      tools: [.functionDeclarations([currentLocationDeclaration, getTemperatureDeclaration])],
+      systemInstruction: ModelContent(parts: """
+      You are a weather bot that specializes in reporting outdoor temperatures in Celsius.
+
+      If not specified, assume that the user is asking for the temperature in their current \
+      location. Use the `getCurrentLocation` function to determine the user's location. Always use the \
+      `getTemperature` function to determine the current outdoor temperature in a location. You \
+      can use the output of the `getCurrentLocation` function as input to the `getTemperature` \
+      function to get the temperature in the user's location.
+
+      Always respond in the format:
+      - Location: City, Province/State, Country
+      - Temperature: #C
+      """)
     )
     let chat = model.startChat()
-    let prompt = """
-    What is the temperature outside right now? Respond in the format:
-    - Location: City, Province/State, Country
-    - Temperature: #C
-
-    Example Output:
-    - Location: Vancouver, British Columbia, Canada
-    - Temperature: 15C
-    """
+    let prompt = "What is the current temperature?"
 
     let response = try await chat.sendMessage(prompt)
 
@@ -340,10 +342,8 @@ struct GenerateContentIntegrationTests {
 
     #expect(response3.functionCalls.isEmpty)
     let finalText = try #require(response3.text).trimmingCharacters(in: .whitespacesAndNewlines)
-    #expect(finalText == """
-    - Location: Waterloo, Ontario, Canada
-    - Temperature: 25C
-    """)
+    #expect(finalText.contains("Waterloo"))
+    #expect(finalText.contains("25"))
 
     if let _ = thinkingConfig.includeThoughts, case .googleAI = config.apiConfig.service {
       #expect(thoughtSignatureCount > 0)

From 323e33b4691513bc3f495563a751deda231eec08 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Thu, 31 Jul 2025 18:46:00 -0400
Subject: [PATCH 18/20] Add `InternalPartTests`

---
 .../Tests/Unit/Types/InternalPartTests.swift  | 86 +++++++++++++++++++
 1 file changed, 86 insertions(+)
 create mode 100644 FirebaseAI/Tests/Unit/Types/InternalPartTests.swift

diff --git a/FirebaseAI/Tests/Unit/Types/InternalPartTests.swift b/FirebaseAI/Tests/Unit/Types/InternalPartTests.swift
new file mode 100644
index 00000000000..3ca72d0fa0f
--- /dev/null
+++ b/FirebaseAI/Tests/Unit/Types/InternalPartTests.swift
@@ -0,0 +1,86 @@
+// Copyright 2024 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.
+
+@testable import FirebaseAI
+import XCTest
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+final class InternalPartTests: XCTestCase {
+  let decoder = JSONDecoder()
+
+  func testDecodeTextPartWithThought() throws {
+    let json = """
+    {
+      "text": "This is a thought.",
+      "thought": true
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+    let part = try decoder.decode(InternalPart.self, from: jsonData)
+
+    XCTAssertEqual(part.isThought, true)
+    guard case let .text(text) = part.data else {
+      XCTFail("Decoded part is not a text part.")
+      return
+    }
+    XCTAssertEqual(text, "This is a thought.")
+  }
+
+  func testDecodeTextPartWithoutThought() throws {
+    let json = """
+    {
+      "text": "This is not a thought."
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+    let part = try decoder.decode(InternalPart.self, from: jsonData)
+
+    XCTAssertNil(part.isThought)
+    guard case let .text(text) = part.data else {
+      XCTFail("Decoded part is not a text part.")
+      return
+    }
+    XCTAssertEqual(text, "This is not a thought.")
+  }
+
+  func testDecodeInlineDataPartWithThought() throws {
+    let imageBase64 =
+      "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+P+/HgAFhAJ/wlseKgAAAABJRU5ErkJggg=="
+    let mimeType = "image/png"
+    let json = """
+    {
+      "inlineData": {
+        "mimeType": "\(mimeType)",
+        "data": "\(imageBase64)"
+      },
+      "thought": true
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+    let part = try decoder.decode(InternalPart.self, from: jsonData)
+
+    XCTAssertEqual(part.isThought, true)
+    guard case let .inlineData(inlineData) = part.data else {
+      XCTFail("Decoded part is not an inlineData part.")
+      return
+    }
+    XCTAssertEqual(inlineData.mimeType, mimeType)
+    XCTAssertEqual(inlineData.data, Data(base64Encoded: imageBase64))
+  }
+
+  // TODO(andrewheard): Add testDecodeInlineDataPartWithoutThought
+  // TODO(andrewheard): Add testDecodeFunctionCallPartWithThought
+  // TODO(andrewheard): Add testDecodeFunctionCallPartWithThoughtSignature
+  // TODO(andrewheard): Add testDecodeFunctionCallPartWithoutThought
+}

From 86f06bec7396dbef294824d9ae62d294b2b6230f Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Fri, 1 Aug 2025 15:01:54 -0400
Subject: [PATCH 19/20] Simplify `generateContentThinkingFunctionCalling` test

---
 .../GenerateContentIntegrationTests.swift     | 59 ++++---------------
 1 file changed, 13 insertions(+), 46 deletions(-)

diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
index b672356a84d..e2900487033 100644
--- a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
+++ b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
@@ -253,11 +253,6 @@ struct GenerateContentIntegrationTests {
   )
   func generateContentThinkingFunctionCalling(_ config: InstanceConfig, modelName: String,
                                               thinkingConfig: ThinkingConfig) async throws {
-    let currentLocationDeclaration = FunctionDeclaration(
-      name: "getCurrentLocation",
-      description: "Returns the user's current city, province or state, and country",
-      parameters: [:]
-    )
     let getTemperatureDeclaration = FunctionDeclaration(
       name: "getTemperature",
       description: "Returns the current temperature in Celsius for the specified location",
@@ -276,15 +271,11 @@ struct GenerateContentIntegrationTests {
         thinkingConfig: thinkingConfig
       ),
       safetySettings: safetySettings,
-      tools: [.functionDeclarations([currentLocationDeclaration, getTemperatureDeclaration])],
+      tools: [.functionDeclarations([getTemperatureDeclaration])],
       systemInstruction: ModelContent(parts: """
       You are a weather bot that specializes in reporting outdoor temperatures in Celsius.
 
-      If not specified, assume that the user is asking for the temperature in their current \
-      location. Use the `getCurrentLocation` function to determine the user's location. Always use the \
-      `getTemperature` function to determine the current outdoor temperature in a location. You \
-      can use the output of the `getCurrentLocation` function as input to the `getTemperature` \
-      function to get the temperature in the user's location.
+      Always use the `getTemperature` function to determine the current temperature in a location.
 
       Always respond in the format:
       - Location: City, Province/State, Country
@@ -292,33 +283,12 @@ struct GenerateContentIntegrationTests {
       """)
     )
     let chat = model.startChat()
-    let prompt = "What is the current temperature?"
+    let prompt = "What is the current temperature in Waterloo, Ontario, Canada?"
 
     let response = try await chat.sendMessage(prompt)
 
-    var thoughtSignatureCount = 0
     #expect(response.functionCalls.count == 1)
-    let locationFunctionCall = try #require(response.functionCalls.first)
-    try #require(locationFunctionCall.name == currentLocationDeclaration.name)
-    #expect(locationFunctionCall.args.isEmpty)
-    #expect(locationFunctionCall.isThought == false)
-    if locationFunctionCall.thoughtSignature != nil {
-      thoughtSignatureCount += 1
-    }
-
-    let locationFunctionResponse = FunctionResponsePart(
-      name: locationFunctionCall.name,
-      response: [
-        "city": .string("Waterloo"),
-        "province": .string("Ontario"),
-        "country": .string("Canada"),
-      ]
-    )
-
-    let response2 = try await chat.sendMessage(locationFunctionResponse)
-
-    #expect(response2.functionCalls.count == 1)
-    let temperatureFunctionCall = try #require(response2.functionCalls.first)
+    let temperatureFunctionCall = try #require(response.functionCalls.first)
     try #require(temperatureFunctionCall.name == getTemperatureDeclaration.name)
     #expect(temperatureFunctionCall.args == [
       "city": .string("Waterloo"),
@@ -326,30 +296,27 @@ struct GenerateContentIntegrationTests {
       "country": .string("Canada"),
     ])
     #expect(temperatureFunctionCall.isThought == false)
-    if temperatureFunctionCall.thoughtSignature != nil {
-      thoughtSignatureCount += 1
+    if let _ = thinkingConfig.includeThoughts, case .googleAI = config.apiConfig.service {
+      let thoughtSignature = try #require(temperatureFunctionCall.thoughtSignature)
+      #expect(!thoughtSignature.isEmpty)
+    } else {
+      #expect(temperatureFunctionCall.thoughtSignature == nil)
     }
 
     let temperatureFunctionResponse = FunctionResponsePart(
-      name: locationFunctionCall.name,
+      name: temperatureFunctionCall.name,
       response: [
         "temperature": .number(25),
         "units": .string("Celsius"),
       ]
     )
 
-    let response3 = try await chat.sendMessage(temperatureFunctionResponse)
+    let response2 = try await chat.sendMessage(temperatureFunctionResponse)
 
-    #expect(response3.functionCalls.isEmpty)
-    let finalText = try #require(response3.text).trimmingCharacters(in: .whitespacesAndNewlines)
+    #expect(response2.functionCalls.isEmpty)
+    let finalText = try #require(response2.text).trimmingCharacters(in: .whitespacesAndNewlines)
     #expect(finalText.contains("Waterloo"))
     #expect(finalText.contains("25"))
-
-    if let _ = thinkingConfig.includeThoughts, case .googleAI = config.apiConfig.service {
-      #expect(thoughtSignatureCount > 0)
-    } else {
-      #expect(thoughtSignatureCount == 0)
-    }
   }
 
   @Test(arguments: [

From 217b06cc7a61e18d0a8535182e7eb468be32a122 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Fri, 1 Aug 2025 15:43:25 -0400
Subject: [PATCH 20/20] Add remaining `InternalPartTests`

---
 .../Tests/Unit/Types/InternalPartTests.swift  | 208 +++++++++++++++++-
 1 file changed, 204 insertions(+), 4 deletions(-)

diff --git a/FirebaseAI/Tests/Unit/Types/InternalPartTests.swift b/FirebaseAI/Tests/Unit/Types/InternalPartTests.swift
index 3ca72d0fa0f..2cd5c5fee2a 100644
--- a/FirebaseAI/Tests/Unit/Types/InternalPartTests.swift
+++ b/FirebaseAI/Tests/Unit/Types/InternalPartTests.swift
@@ -27,6 +27,7 @@ final class InternalPartTests: XCTestCase {
     }
     """
     let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
     let part = try decoder.decode(InternalPart.self, from: jsonData)
 
     XCTAssertEqual(part.isThought, true)
@@ -44,6 +45,7 @@ final class InternalPartTests: XCTestCase {
     }
     """
     let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
     let part = try decoder.decode(InternalPart.self, from: jsonData)
 
     XCTAssertNil(part.isThought)
@@ -68,6 +70,7 @@ final class InternalPartTests: XCTestCase {
     }
     """
     let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
     let part = try decoder.decode(InternalPart.self, from: jsonData)
 
     XCTAssertEqual(part.isThought, true)
@@ -79,8 +82,205 @@ final class InternalPartTests: XCTestCase {
     XCTAssertEqual(inlineData.data, Data(base64Encoded: imageBase64))
   }
 
-  // TODO(andrewheard): Add testDecodeInlineDataPartWithoutThought
-  // TODO(andrewheard): Add testDecodeFunctionCallPartWithThought
-  // TODO(andrewheard): Add testDecodeFunctionCallPartWithThoughtSignature
-  // TODO(andrewheard): Add testDecodeFunctionCallPartWithoutThought
+  func testDecodeInlineDataPartWithoutThought() throws {
+    let imageBase64 = "aGVsbG8="
+    let mimeType = "image/png"
+    let json = """
+    {
+      "inlineData": {
+        "mimeType": "\(mimeType)",
+        "data": "\(imageBase64)"
+      }
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let part = try decoder.decode(InternalPart.self, from: jsonData)
+
+    XCTAssertNil(part.isThought)
+    guard case let .inlineData(inlineData) = part.data else {
+      XCTFail("Decoded part is not an inlineData part.")
+      return
+    }
+    XCTAssertEqual(inlineData.mimeType, mimeType)
+    XCTAssertEqual(inlineData.data, Data(base64Encoded: imageBase64))
+  }
+
+  func testDecodeFileDataPartWithThought() throws {
+    let uri = "file:///path/to/file.mp3"
+    let mimeType = "audio/mpeg"
+    let json = """
+    {
+      "fileData": {
+        "fileUri": "\(uri)",
+        "mimeType": "\(mimeType)"
+      },
+      "thought": true
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let part = try decoder.decode(InternalPart.self, from: jsonData)
+
+    XCTAssertEqual(part.isThought, true)
+    guard case let .fileData(fileData) = part.data else {
+      XCTFail("Decoded part is not a fileData part.")
+      return
+    }
+    XCTAssertEqual(fileData.fileURI, uri)
+    XCTAssertEqual(fileData.mimeType, mimeType)
+  }
+
+  func testDecodeFileDataPartWithoutThought() throws {
+    let uri = "file:///path/to/file.mp3"
+    let mimeType = "audio/mpeg"
+    let json = """
+    {
+      "fileData": {
+        "fileUri": "\(uri)",
+        "mimeType": "\(mimeType)"
+      }
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let part = try decoder.decode(InternalPart.self, from: jsonData)
+
+    XCTAssertNil(part.isThought)
+    guard case let .fileData(fileData) = part.data else {
+      XCTFail("Decoded part is not a fileData part.")
+      return
+    }
+    XCTAssertEqual(fileData.fileURI, uri)
+    XCTAssertEqual(fileData.mimeType, mimeType)
+  }
+
+  func testDecodeFunctionCallPartWithThoughtSignature() throws {
+    let functionName = "someFunction"
+    let expectedThoughtSignature = "some_signature"
+    let json = """
+    {
+      "functionCall": {
+        "name": "\(functionName)",
+        "args": {
+          "arg1": "value1"
+        },
+      },
+      "thoughtSignature": "\(expectedThoughtSignature)"
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let part = try decoder.decode(InternalPart.self, from: jsonData)
+
+    let thoughtSignature = try XCTUnwrap(part.thoughtSignature)
+    XCTAssertEqual(thoughtSignature, expectedThoughtSignature)
+    XCTAssertNil(part.isThought)
+    guard case let .functionCall(functionCall) = part.data else {
+      XCTFail("Decoded part is not a functionCall part.")
+      return
+    }
+    XCTAssertEqual(functionCall.name, functionName)
+    XCTAssertEqual(functionCall.args, ["arg1": .string("value1")])
+  }
+
+  func testDecodeFunctionCallPartWithoutThoughtSignature() throws {
+    let functionName = "someFunction"
+    let json = """
+    {
+      "functionCall": {
+        "name": "\(functionName)",
+        "args": {
+          "arg1": "value1"
+        }
+      }
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let part = try decoder.decode(InternalPart.self, from: jsonData)
+
+    XCTAssertNil(part.isThought)
+    XCTAssertNil(part.thoughtSignature)
+    guard case let .functionCall(functionCall) = part.data else {
+      XCTFail("Decoded part is not a functionCall part.")
+      return
+    }
+    XCTAssertEqual(functionCall.name, functionName)
+    XCTAssertEqual(functionCall.args, ["arg1": .string("value1")])
+  }
+
+  func testDecodeFunctionCallPartWithoutArgs() throws {
+    let functionName = "someFunction"
+    let json = """
+    {
+      "functionCall": {
+        "name": "\(functionName)"
+      }
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let part = try decoder.decode(InternalPart.self, from: jsonData)
+
+    XCTAssertNil(part.isThought)
+    XCTAssertNil(part.thoughtSignature)
+    guard case let .functionCall(functionCall) = part.data else {
+      XCTFail("Decoded part is not a functionCall part.")
+      return
+    }
+    XCTAssertEqual(functionCall.name, functionName)
+    XCTAssertEqual(functionCall.args, JSONObject())
+  }
+
+  func testDecodeFunctionResponsePartWithThought() throws {
+    let functionName = "someFunction"
+    let json = """
+    {
+      "functionResponse": {
+        "name": "\(functionName)",
+        "response": {
+          "output": "someValue"
+        }
+      },
+      "thought": true
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let part = try decoder.decode(InternalPart.self, from: jsonData)
+
+    XCTAssertEqual(part.isThought, true)
+    guard case let .functionResponse(functionResponse) = part.data else {
+      XCTFail("Decoded part is not a functionResponse part.")
+      return
+    }
+    XCTAssertEqual(functionResponse.name, functionName)
+    XCTAssertEqual(functionResponse.response, ["output": .string("someValue")])
+  }
+
+  func testDecodeFunctionResponsePartWithoutThought() throws {
+    let functionName = "someFunction"
+    let json = """
+    {
+      "functionResponse": {
+        "name": "\(functionName)",
+        "response": {
+          "output": "someValue"
+        }
+      }
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let part = try decoder.decode(InternalPart.self, from: jsonData)
+
+    XCTAssertNil(part.isThought)
+    guard case let .functionResponse(functionResponse) = part.data else {
+      XCTFail("Decoded part is not a functionResponse part.")
+      return
+    }
+    XCTAssertEqual(functionResponse.name, functionName)
+    XCTAssertEqual(functionResponse.response, ["output": .string("someValue")])
+  }
 }