feat(firebaseai): live session resumption

packages/firebase_ai/firebase_ai/example/lib/pages/chat_page.dart

@@ -153,6 +153,7 @@ class _ChatPageState extends State<ChatPage> {
                             Icons.send,
                             color: Theme.of(context).colorScheme.primary,
                           ),
+                          tooltip: 'Send',
                         ),
                         IconButton(
                           onPressed: () {
@@ -162,6 +163,7 @@ class _ChatPageState extends State<ChatPage> {
                             Icons.stream,
                             color: Theme.of(context).colorScheme.primary,
                           ),
+                          tooltip: 'Send Stream',
                         ),
                       ],
                     )

packages/firebase_ai/firebase_ai/example/lib/utils/audio_input.dart

@@ -110,8 +110,6 @@ class AudioInput extends ChangeNotifier {
       sampleRate: 24000,
       device: selectedDevice,
       numChannels: 1,
-      echoCancel: true,
-      noiseSuppress: true,
       androidConfig: const AndroidRecordConfig(
         audioSource: AndroidAudioSource.voiceCommunication,
       ),

packages/firebase_ai/firebase_ai/lib/firebase_ai.dart

@@ -102,15 +102,19 @@ export 'src/imagen/imagen_reference.dart'
         ImagenControlReference;
 export 'src/live_api.dart'
     show
-        LiveGenerationConfig,
-        SpeechConfig,
         AudioTranscriptionConfig,
+        ContextWindowCompressionConfig,
+        GoingAwayNotice,
+        LiveGenerationConfig,
         LiveServerMessage,
         LiveServerContent,
         LiveServerToolCall,
         LiveServerToolCallCancellation,
         LiveServerResponse,
-        GoingAwayNotice,
+        SessionResumptionConfig,
+        SessionResumptionUpdate,
+        SlidingWindow,
+        SpeechConfig,
         Transcription;
 export 'src/live_session.dart' show LiveSession;
 export 'src/schema.dart' show JSONSchema, Schema, SchemaType;

packages/firebase_ai/firebase_ai/lib/src/base_model.dart

@@ -13,16 +13,12 @@
 // limitations under the License.
 
 import 'dart:async';
-import 'dart:convert';
 
 import 'package:firebase_app_check/firebase_app_check.dart';
 import 'package:firebase_auth/firebase_auth.dart';
 import 'package:firebase_core/firebase_core.dart';
-import 'package:flutter/foundation.dart';
 import 'package:http/http.dart' as http;
 import 'package:meta/meta.dart';
-import 'package:web_socket_channel/io.dart';
-import 'package:web_socket_channel/web_socket_channel.dart';
 
 import 'api.dart';
 import 'client.dart';

packages/firebase_ai/firebase_ai/lib/src/live_api.dart

@@ -77,21 +77,90 @@ class AudioTranscriptionConfig {
   Map<String, Object?> toJson() => {};
 }
 
+/// Configures the sliding window context compression mechanism.
+///
+/// The context window will be truncated by keeping only a suffix of it.
+class SlidingWindow {
+  /// Creates a [SlidingWindow] instance.
+  ///
+  /// [targetTokens] (optional): The target number of tokens to keep in the
+  /// context window.
+  SlidingWindow({this.targetTokens});
+
+  /// The session reduction target, i.e., how many tokens we should keep.
+  final int? targetTokens;
+  // ignore: public_member_api_docs
+  Map<String, Object?> toJson() =>
+      {if (targetTokens case final targetTokens?) 'targetTokens': targetTokens};
+}
+
+/// Enables context window compression to manage the model's context window.
+///
+/// This mechanism prevents the context from exceeding a given length.
+class ContextWindowCompressionConfig {
+  /// Creates a [ContextWindowCompressionConfig] instance.
+  ///
+  /// [triggerTokens] (optional): The number of tokens that triggers the
+  /// compression mechanism.
+  /// [slidingWindow] (optional): The sliding window compression mechanism to
+  /// use.
+  ContextWindowCompressionConfig({this.triggerTokens, this.slidingWindow});
+
+  /// The number of tokens (before running a turn) that triggers the context
+  /// window compression.
+  final int? triggerTokens;
+
+  /// The sliding window compression mechanism.
+  final SlidingWindow? slidingWindow;
+  // ignore: public_member_api_docs
+  Map<String, Object?> toJson() => {
+        if (triggerTokens case final triggerTokens?)
+          'triggerTokens': triggerTokens,
+        if (slidingWindow case final slidingWindow?)
+          'slidingWindow': slidingWindow.toJson()
+      };
+}
+
+/// Configuration for the session resumption mechanism.
+///
+/// When included in the session setup, the server will send
+/// [SessionResumptionUpdate] messages.
+class SessionResumptionConfig {
+  /// Creates a [SessionResumptionConfig] instance.
+  ///
+  /// [handle] (optional): The session resumption handle of the previous session
+  /// to restore.
+  /// [transparent] (optional): If set, requests the server to send updates with
+  /// the message index of the last client message included in the session
+  /// state.
+  SessionResumptionConfig({this.handle});
+
+  /// The session resumption handle of the previous session to restore.
+  ///
+  /// If not present, a new session will be started.
+  final String? handle;
+
+  // ignore: public_member_api_docs
+  Map<String, Object?> toJson() => {
+        if (handle case final handle?) 'handle': handle,
+      };
+}
+
 /// Configures live generation settings.
 final class LiveGenerationConfig extends BaseGenerationConfig {
   // ignore: public_member_api_docs
-  LiveGenerationConfig({
-    this.speechConfig,
-    this.inputAudioTranscription,
-    this.outputAudioTranscription,
-    super.responseModalities,
-    super.maxOutputTokens,
-    super.temperature,
-    super.topP,
-    super.topK,
-    super.presencePenalty,
-    super.frequencyPenalty,
-  });
+  LiveGenerationConfig(
+      {this.speechConfig,
+      this.inputAudioTranscription,
+      this.outputAudioTranscription,
+      this.contextWindowCompression,
+      super.responseModalities,
+      super.maxOutputTokens,
+      super.temperature,
+      super.topP,
+      super.topK,
+      super.presencePenalty,
+      super.frequencyPenalty});
 
   /// The speech configuration.
   final SpeechConfig? speechConfig;
@@ -103,6 +172,9 @@ final class LiveGenerationConfig extends BaseGenerationConfig {
   /// the output audio.
   final AudioTranscriptionConfig? outputAudioTranscription;
 
+  /// The context window compression configuration.
+  final ContextWindowCompressionConfig? contextWindowCompression;
+
   @override
   Map<String, Object?> toJson() => {
         ...super.toJson(),
@@ -222,6 +294,34 @@ class GoingAwayNotice implements LiveServerMessage {
   final String? timeLeft;
 }
 
+/// An update of the session resumption state.
+///
+/// This message is only sent if [SessionResumptionConfig] was set in the
+/// session setup.
+class SessionResumptionUpdate implements LiveServerMessage {
+  /// Creates a [SessionResumptionUpdate] instance.
+  ///
+  /// [newHandle] (optional): The new handle that represents the state that can
+  /// be resumed.
+  /// [resumable] (optional): Indicates if the session can be resumed at this
+  /// point.
+  /// [lastConsumedClientMessageIndex] (optional): The index of the last client
+  /// message that is included in the state represented by this update.
+  SessionResumptionUpdate(
+      {this.newHandle, this.resumable, this.lastConsumedClientMessageIndex});
+
+  /// The new handle that represents the state that can be resumed. Empty if
+  /// `resumable` is false.
+  final String? newHandle;
+
+  /// Indicates if the session can be resumed at this point.
+  final bool? resumable;
+
+  /// The index of the last client message that is included in the state
+  /// represented by this update.
+  final int? lastConsumedClientMessageIndex;
+}
+
 /// A single response chunk received during a live content generation.
 ///
 /// It can contain generated content, function calls to be executed, or
@@ -449,8 +549,17 @@ LiveServerMessage _parseServerMessage(Object jsonObject) {
   } else if (json.containsKey('setupComplete')) {
     return LiveServerSetupComplete();
   } else if (json.containsKey('goAway')) {
-    final goAwayJson = json['goAway'] as Map;
+    final goAwayJson = json['goAway'] as Map<String, dynamic>;
     return GoingAwayNotice(timeLeft: goAwayJson['timeLeft'] as String?);
+  } else if (json.containsKey('sessionResumptionUpdate')) {
+    final sessionResumptionUpdateJson =
+        json['sessionResumptionUpdate'] as Map<String, dynamic>;
+    return SessionResumptionUpdate(
+      newHandle: sessionResumptionUpdateJson['newHandle'] as String?,
+      resumable: sessionResumptionUpdateJson['resumable'] as bool?,
+      lastConsumedClientMessageIndex:
+          sessionResumptionUpdateJson['lastConsumedClientMessageIndex'] as int?,
+    );
   } else {
     throw unhandledFormat('LiveServerMessage', json);
   }

packages/firebase_ai/firebase_ai/lib/src/live_model.dart

@@ -93,45 +93,28 @@ final class LiveGenerativeModel extends BaseModel {
   ///
   /// Returns a [Future] that resolves to an [LiveSession] object upon successful
   /// connection.
-  Future<LiveSession> connect() async {
+  Future<LiveSession> connect(
+      {SessionResumptionConfig? sessionResumption}) async {
     final uri = _useVertexBackend ? _vertexAIUri() : _googleAIUri();
     final modelString =
         _useVertexBackend ? _vertexAIModelString() : _googleAIModelString();
 
-    final setupJson = {
-      'setup': {
-        'model': modelString,
-        if (_systemInstruction != null)
-          'system_instruction': _systemInstruction.toJson(),
-        if (_tools != null) 'tools': _tools.map((t) => t.toJson()).toList(),
-        if (_liveGenerationConfig != null) ...{
-          'generation_config': _liveGenerationConfig.toJson(),
-          if (_liveGenerationConfig.inputAudioTranscription != null)
-            'input_audio_transcription':
-                _liveGenerationConfig.inputAudioTranscription!.toJson(),
-          if (_liveGenerationConfig.outputAudioTranscription != null)
-            'output_audio_transcription':
-                _liveGenerationConfig.outputAudioTranscription!.toJson(),
-        },
-      }
-    };
-
-    final request = jsonEncode(setupJson);
     final headers = await BaseModel.firebaseTokens(
       _appCheck,
       _auth,
       _app,
       _useLimitedUseAppCheckTokens,
     )();
 
-    var ws = kIsWeb
-        ? WebSocketChannel.connect(Uri.parse(uri))
-        : IOWebSocketChannel.connect(Uri.parse(uri), headers: headers);
-    await ws.ready;
-
-    ws.sink.add(request);
-
-    return LiveSession(ws);
+    return LiveSession.connect(
+      uri: uri,
+      headers: headers,
+      modelString: modelString,
+      systemInstruction: _systemInstruction,
+      tools: _tools,
+      sessionResumption: sessionResumption,
+      liveGenerationConfig: _liveGenerationConfig,
+    );
   }
 }