diff --git a/packages/firebase_ai/firebase_ai/lib/src/api.dart b/packages/firebase_ai/firebase_ai/lib/src/api.dart
index 0c1849c3c18d..ab90340b6a29 100644
--- a/packages/firebase_ai/firebase_ai/lib/src/api.dart
+++ b/packages/firebase_ai/firebase_ai/lib/src/api.dart
@@ -259,7 +259,7 @@ final class Candidate {
   ///
   /// If this candidate was finished for a reason of [FinishReason.recitation]
   /// or [FinishReason.safety], accessing this text will throw a
-  /// [GenerativeAIException].
+  /// [FirebaseAIException].
   ///
   /// If [content] contains any text parts, this value is the concatenation of
   /// the text.
diff --git a/packages/firebase_ai/firebase_ai/lib/src/chat.dart b/packages/firebase_ai/firebase_ai/lib/src/chat.dart
index 6f6d32d6f6f0..6a0846fd0214 100644
--- a/packages/firebase_ai/firebase_ai/lib/src/chat.dart
+++ b/packages/firebase_ai/firebase_ai/lib/src/chat.dart
@@ -1,16 +1,16 @@
-// // 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.
+// 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.
 
 import 'dart:async';
 
diff --git a/packages/firebase_ai/firebase_ai/lib/src/content.dart b/packages/firebase_ai/firebase_ai/lib/src/content.dart
index fac91c3650a0..c9ed7813c4ee 100644
--- a/packages/firebase_ai/firebase_ai/lib/src/content.dart
+++ b/packages/firebase_ai/firebase_ai/lib/src/content.dart
@@ -16,6 +16,8 @@ import 'dart:convert';
 import 'dart:developer';
 import 'dart:typed_data';
 
+import 'package:meta/meta.dart';
+
 import 'api.dart';
 import 'error.dart';
 
@@ -120,9 +122,11 @@ Part parsePart(Object? jsonObject) {
     if (executableCode is Map &&
         executableCode.containsKey('language') &&
         executableCode.containsKey('code')) {
-      return ExecutableCodePart(
+      return ExecutableCodePart._(
         language: CodeLanguage.parseValue(executableCode['language'] as String),
         code: executableCode['code'] as String,
+        isThought: isThought,
+        thoughtSignature: thoughtSignature,
       );
     } else {
       throw unhandledFormat('executableCode', executableCode);
@@ -133,9 +137,11 @@ Part parsePart(Object? jsonObject) {
     if (codeExecutionResult is Map &&
         codeExecutionResult.containsKey('outcome') &&
         codeExecutionResult.containsKey('output')) {
-      return CodeExecutionResultPart(
+      return CodeExecutionResultPart._(
         outcome: Outcome.parseValue(codeExecutionResult['outcome'] as String),
         output: codeExecutionResult['output'] as String,
+        isThought: isThought,
+        thoughtSignature: thoughtSignature,
       );
     } else {
       throw unhandledFormat('codeExecutionResult', codeExecutionResult);
@@ -188,7 +194,11 @@ sealed class Part {
   final String? _thoughtSignature;
 
   /// Convert the [Part] content to json format.
-  Object toJson();
+  Object toJson() => {
+        if (isThought case final isThought?) 'thought': isThought,
+        if (_thoughtSignature case final thoughtSignature?)
+          'thoughtSignature': thoughtSignature,
+      };
 }
 
 /// A [Part] that contains unparsable data.
@@ -200,7 +210,10 @@ final class UnknownPart extends Part {
   final Map<String, Object?> data;
 
   @override
-  Object toJson() => data;
+  Object toJson() {
+    final superJson = super.toJson() as Map<String, Object?>;
+    return <String, Object?>{...superJson, ...data};
+  }
 }
 
 /// A [Part] with the text content.
@@ -211,6 +224,18 @@ final class TextPart extends Part {
           isThought: isThought,
           thoughtSignature: null,
         );
+
+  @visibleForTesting
+  // ignore: public_member_api_docs
+  const TextPart.forTest(
+    this.text, {
+    bool? isThought,
+    String? thoughtSignature,
+  }) : super(
+          isThought: isThought,
+          thoughtSignature: thoughtSignature,
+        );
+
   const TextPart._(
     this.text, {
     bool? isThought,
@@ -223,7 +248,10 @@ final class TextPart extends Part {
   /// The text content of the [Part]
   final String text;
   @override
-  Object toJson() => {'text': text};
+  Object toJson() {
+    final superJson = super.toJson() as Map<String, Object?>;
+    return <String, Object?>{...superJson, 'text': text};
+  }
 }
 
 /// A [Part] with the byte content of a file.
@@ -238,6 +266,20 @@ final class InlineDataPart extends Part {
           isThought: isThought,
           thoughtSignature: null,
         );
+
+  @visibleForTesting
+  // ignore: public_member_api_docs
+  const InlineDataPart.forTest(
+    this.mimeType,
+    this.bytes, {
+    this.willContinue,
+    bool? isThought,
+    String? thoughtSignature,
+  }) : super(
+          isThought: isThought,
+          thoughtSignature: thoughtSignature,
+        );
+
   const InlineDataPart._(
     this.mimeType,
     this.bytes, {
@@ -259,13 +301,17 @@ final class InlineDataPart extends Part {
   /// Whether there's more inline data coming for streaming.
   final bool? willContinue;
   @override
-  Object toJson() => {
-        'inlineData': {
-          'data': base64Encode(bytes),
-          'mimeType': mimeType,
-          if (willContinue != null) 'willContinue': willContinue,
-        }
-      };
+  Object toJson() {
+    final superJson = super.toJson() as Map<String, Object?>;
+    return <String, Object?>{
+      ...superJson,
+      'inlineData': {
+        'data': base64Encode(bytes),
+        'mimeType': mimeType,
+        if (willContinue != null) 'willContinue': willContinue,
+      },
+    };
+  }
 
   /// The representation of the data in media streaming chunk.
   Object toMediaChunkJson() => {
@@ -289,6 +335,20 @@ final class FunctionCall extends Part {
           isThought: isThought,
           thoughtSignature: null,
         );
+
+  @visibleForTesting
+  // ignore: public_member_api_docs
+  const FunctionCall.forTest(
+    this.name,
+    this.args, {
+    this.id,
+    bool? isThought,
+    String? thoughtSignature,
+  }) : super(
+          isThought: isThought,
+          thoughtSignature: thoughtSignature,
+        );
+
   const FunctionCall._(
     this.name,
     this.args, {
@@ -313,13 +373,17 @@ final class FunctionCall extends Part {
   final String? id;
 
   @override
-  Object toJson() => {
-        'functionCall': {
-          'name': name,
-          'args': args,
-          if (id != null) 'id': id,
-        }
-      };
+  Object toJson() {
+    final superJson = super.toJson() as Map<String, Object?>;
+    return <String, Object?>{
+      ...superJson,
+      'functionCall': {
+        'name': name,
+        'args': args,
+        if (id != null) 'id': id,
+      },
+    };
+  }
 }
 
 /// The response class for [FunctionCall]
@@ -350,13 +414,17 @@ final class FunctionResponse extends Part {
   final String? id;
 
   @override
-  Object toJson() => {
-        'functionResponse': {
-          'name': name,
-          'response': response,
-          if (id != null) 'id': id,
-        }
-      };
+  Object toJson() {
+    final superJson = super.toJson() as Map<String, Object?>;
+    return <String, Object?>{
+      ...superJson,
+      'functionResponse': {
+        'name': name,
+        'response': response,
+        if (id != null) 'id': id,
+      },
+    };
+  }
 }
 
 /// A [Part] with Firebase Storage uri as prompt content
@@ -370,6 +438,19 @@ final class FileData extends Part {
           isThought: isThought,
           thoughtSignature: null,
         );
+
+  @visibleForTesting
+  // ignore: public_member_api_docs
+  const FileData.forTest(
+    this.mimeType,
+    this.fileUri, {
+    bool? isThought,
+    String? thoughtSignature,
+  }) : super(
+          isThought: isThought,
+          thoughtSignature: thoughtSignature,
+        );
+
   const FileData._(
     this.mimeType,
     this.fileUri, {
@@ -388,9 +469,13 @@ final class FileData extends Part {
   final String fileUri;
 
   @override
-  Object toJson() => {
-        'file_data': {'file_uri': fileUri, 'mime_type': mimeType}
-      };
+  Object toJson() {
+    final superJson = super.toJson() as Map<String, Object?>;
+    return <String, Object?>{
+      ...superJson,
+      'file_data': {'file_uri': fileUri, 'mime_type': mimeType},
+    };
+  }
 }
 
 /// A `Part` that represents the code that is executed by the model.
@@ -405,6 +490,16 @@ final class ExecutableCodePart extends Part {
           thoughtSignature: null,
         );
 
+  ExecutableCodePart._({
+    required this.language,
+    required this.code,
+    bool? isThought,
+    String? thoughtSignature,
+  }) : super(
+          isThought: isThought,
+          thoughtSignature: thoughtSignature,
+        );
+
   /// The programming language of the code.
   final CodeLanguage language;
 
@@ -412,9 +507,13 @@ final class ExecutableCodePart extends Part {
   final String code;
 
   @override
-  Object toJson() => {
-        'executableCode': {'language': language.toJson(), 'code': code}
-      };
+  Object toJson() {
+    final superJson = super.toJson() as Map<String, Object?>;
+    return <String, Object?>{
+      ...superJson,
+      'executableCode': {'language': language.toJson(), 'code': code},
+    };
+  }
 }
 
 /// A `Part` that represents the code execution result from the model.
@@ -429,6 +528,16 @@ final class CodeExecutionResultPart extends Part {
           thoughtSignature: null,
         );
 
+  CodeExecutionResultPart._({
+    required this.outcome,
+    required this.output,
+    bool? isThought,
+    String? thoughtSignature,
+  }) : super(
+          isThought: isThought,
+          thoughtSignature: thoughtSignature,
+        );
+
   /// The result of the execution.
   final Outcome outcome;
 
@@ -436,7 +545,11 @@ final class CodeExecutionResultPart extends Part {
   final String output;
 
   @override
-  Object toJson() => {
-        'codeExecutionResult': {'outcome': outcome.toJson(), 'output': output}
-      };
+  Object toJson() {
+    final superJson = super.toJson() as Map<String, Object?>;
+    return <String, Object?>{
+      ...superJson,
+      'codeExecutionResult': {'outcome': outcome.toJson(), 'output': output},
+    };
+  }
 }
diff --git a/packages/firebase_ai/firebase_ai/test/content_test.dart b/packages/firebase_ai/firebase_ai/test/content_test.dart
index 6b7255f0d5db..d9c50082b1dd 100644
--- a/packages/firebase_ai/firebase_ai/test/content_test.dart
+++ b/packages/firebase_ai/firebase_ai/test/content_test.dart
@@ -82,81 +82,105 @@ void main() {
   });
 
   group('Part tests', () {
-    test('TextPart toJson', () {
-      const part = TextPart('Test');
-      final json = part.toJson();
-      expect((json as Map)['text'], 'Test');
+    test('TextPart with isThought and thoughtSignature toJson', () {
+      const part =
+          TextPart.forTest('Test', isThought: true, thoughtSignature: 'sig');
+      final json = part.toJson() as Map;
+      expect(json['text'], 'Test');
+      expect(json['thought'], true);
+      expect(json['thoughtSignature'], 'sig');
     });
 
-    test('DataPart toJson', () {
-      final part = InlineDataPart('image/png', Uint8List(0));
-      final json = part.toJson();
-      expect((json as Map)['inlineData']['mimeType'], 'image/png');
-      expect(json['inlineData']['data'], '');
-      expect(json['inlineData'].containsKey('willContinue'), false);
+    test('DataPart with isThought and thoughtSignature toJson', () {
+      final part = InlineDataPart.forTest('image/png', Uint8List(0),
+          isThought: true, thoughtSignature: 'sig');
+      final json = part.toJson() as Map;
+      final inlineData = json['inlineData'] as Map;
+      expect(inlineData['mimeType'], 'image/png');
+      expect(inlineData['data'], '');
+      expect(json.containsKey('willContinue'), false);
+      expect(json['thought'], true);
+      expect(json['thoughtSignature'], 'sig');
     });
 
     test('DataPart with false willContinue toJson', () {
       final part =
           InlineDataPart('image/png', Uint8List(0), willContinue: false);
-      final json = part.toJson();
-      expect((json as Map)['inlineData']['mimeType'], 'image/png');
-      expect(json['inlineData']['data'], '');
-      expect(json['inlineData'].containsKey('willContinue'), true);
-      expect(json['inlineData']['willContinue'], false);
+      final json = part.toJson() as Map;
+      final inlineData = json['inlineData'] as Map;
+      expect(inlineData['mimeType'], 'image/png');
+      expect(inlineData['data'], '');
+      expect(inlineData.containsKey('willContinue'), true);
+      expect(inlineData['willContinue'], false);
     });
 
     test('DataPart with true willContinue toJson', () {
       final part =
           InlineDataPart('image/png', Uint8List(0), willContinue: true);
-      final json = part.toJson();
-      expect((json as Map)['inlineData']['mimeType'], 'image/png');
-      expect(json['inlineData']['data'], '');
-      expect(json['inlineData'].containsKey('willContinue'), true);
-      expect(json['inlineData']['willContinue'], true);
+      final json = part.toJson() as Map;
+      final inlineData = json['inlineData'] as Map;
+      expect(inlineData['mimeType'], 'image/png');
+      expect(inlineData['data'], '');
+      expect(inlineData.containsKey('willContinue'), true);
+      expect(inlineData['willContinue'], true);
     });
 
-    test('FunctionCall toJson', () {
-      const part = FunctionCall(
+    test('FunctionCall with isThought and thoughtSignature toJson', () {
+      const part = FunctionCall.forTest(
           'myFunction',
           {
             'arguments': [
               {'text': 'Test'}
             ],
           },
-          id: 'myFunctionId');
-      final json = part.toJson();
-      expect((json as Map)['functionCall']['name'], 'myFunction');
-      expect(json['functionCall']['args'].length, 1);
-      expect(json['functionCall']['args']['arguments'].length, 1);
-      expect(json['functionCall']['args']['arguments'][0]['text'], 'Test');
-      expect(json['functionCall']['id'], 'myFunctionId');
+          id: 'myFunctionId',
+          isThought: true,
+          thoughtSignature: 'sig');
+      final json = part.toJson() as Map;
+      final functionCall = json['functionCall'] as Map;
+      expect(functionCall['name'], 'myFunction');
+      final args = functionCall['args'] as Map;
+      expect(args.length, 1);
+      final arguments = args['arguments'] as List;
+      expect(arguments.length, 1);
+      final text = arguments[0] as Map;
+      expect(text['text'], 'Test');
+      expect(functionCall['id'], 'myFunctionId');
+      expect(json['thought'], true);
+      expect(json['thoughtSignature'], 'sig');
     });
 
-    test('FunctionResponse toJson', () {
+    test('FunctionResponse with isThought', () {
       final part = FunctionResponse(
-          'myFunction',
-          {
-            'inlineData': {
-              'mimeType': 'application/octet-stream',
-              'data': Uint8List(0)
-            }
-          },
-          id: 'myFunctionId');
-      final json = part.toJson();
-      expect((json as Map)['functionResponse']['name'], 'myFunction');
-      expect(json['functionResponse']['response']['inlineData']['mimeType'],
-          'application/octet-stream');
-      expect(json['functionResponse']['response']['inlineData']['data'],
-          Uint8List(0));
-      expect(json['functionResponse']['id'], 'myFunctionId');
+        'myFunction',
+        {
+          'inlineData': {
+            'mimeType': 'application/octet-stream',
+            'data': Uint8List(0)
+          }
+        },
+        id: 'myFunctionId',
+        isThought: true,
+      );
+      final json = part.toJson() as Map;
+      final functionResponse = json['functionResponse'] as Map;
+      expect(functionResponse['name'], 'myFunction');
+      final response = functionResponse['response'] as Map;
+      final inlineData = response['inlineData'] as Map;
+      expect(inlineData['mimeType'], 'application/octet-stream');
+      expect(inlineData['data'], Uint8List(0));
+      expect(functionResponse['id'], 'myFunctionId');
+      expect(json['thought'], true);
     });
 
-    test('FileData toJson', () {
-      const part = FileData('image/png', 'gs://bucket-name/path');
-      final json = part.toJson();
-      expect((json as Map)['file_data']['mime_type'], 'image/png');
-      expect(json['file_data']['file_uri'], 'gs://bucket-name/path');
+    test('FileData with isThought and thoughtSignature toJson', () {
+      const part = FileData.forTest('image/png', 'gs://bucket-name/path',
+          isThought: true);
+      final json = part.toJson() as Map;
+      final fileData = json['file_data'] as Map;
+      expect(fileData['mime_type'], 'image/png');
+      expect(fileData['file_uri'], 'gs://bucket-name/path');
+      expect(json['thought'], true);
     });
   });