Skip to content

FirebaseAI: Implement live session resumption and context window compression#1410

Draft
cynthiajoan wants to merge 1 commit intomainfrom
firebaseai/live-session-resumption
Draft

FirebaseAI: Implement live session resumption and context window compression#1410
cynthiajoan wants to merge 1 commit intomainfrom
firebaseai/live-session-resumption

Conversation

@cynthiajoan
Copy link
Contributor

…n with new configuration classes and updated connection logic.

Description

Provide details of the change, and generalize the change in the PR title above.


Testing

Describe how you've tested these changes.


Type of Change

Place an x the applicable box:

  • Bug fix. Add the issue # below if applicable.
  • New feature. A non-breaking change which adds functionality.
  • Other, such as a build process or documentation change.

…n with new configuration classes and updated connection logic.
@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the Firebase AI SDK by introducing advanced session management capabilities. It enables live AI sessions to be resumed, preventing loss of conversational context due to connection interruptions. Additionally, it integrates context window compression, allowing for more efficient use of the model's context window by intelligently reducing session length based on configurable triggers and targets. These changes provide a more resilient and optimized experience for long-running AI interactions.

Highlights

  • Live Session Resumption: Implemented the ability to resume live AI sessions, allowing for seamless re-establishment of connections without losing context. This includes a new SessionResumptionConfig class and a ResumeSessionAsync method in LiveSession.
  • Context Window Compression: Introduced context window compression to manage the model's context more efficiently. New configuration classes, ContextWindowCompressionConfig and SlidingWindow, were added to define compression triggers and targets.
  • Updated Connection Logic: Refactored the LiveGenerativeModel.ConnectAsync method to support session resumption and context window compression, making the connection process more robust and configurable.
  • Dynamic WebSocket Handling: Modified LiveSession to dynamically manage its underlying ClientWebSocket instance, enabling the replacement of the WebSocket during session resumption while ensuring thread safety during message reception.
  • Session Resumption Updates: Added a new LiveSessionResumptionUpdate message type to LiveSessionResponse to communicate the state of session resumption, including new handles and resumability status.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog
  • firebaseai/src/ContextWindowCompressionConfig.cs
    • Added SlidingWindow class for configuring token reduction in context compression.
    • Added ContextWindowCompressionConfig class to enable and configure context window compression.
  • firebaseai/src/LiveGenerationConfig.cs
    • Added _contextWindowCompression field and ContextWindowCompression property.
    • Updated constructor to accept ContextWindowCompressionConfig.
    • Modified ToJson method to include context window compression configuration.
  • firebaseai/src/LiveGenerativeModel.cs
    • Modified ConnectAsync to accept SessionResumptionConfig.
    • Refactored connection setup into a factory function to facilitate session resumption.
    • Included session resumption and context window compression configurations in the initial setup message.
    • Updated ConnectAsync to return LiveSession with the connection factory for future resumptions.
  • firebaseai/src/LiveSession.cs
    • Changed _clientWebSocket to be mutable and added _connectionFactory for re-establishing connections.
    • Updated constructor to accept the connection factory.
    • Implemented ResumeSessionAsync method to close the old WebSocket and establish a new one.
    • Modified ReceiveAsync to safely handle changes to the underlying ClientWebSocket during reception, particularly during session resumption.
  • firebaseai/src/LiveSessionResponse.cs
    • Updated FromJson method to parse sessionResumptionUpdate messages.
    • Added LiveSessionResumptionUpdate struct to represent session resumption state updates.
  • firebaseai/src/SessionResumptionConfig.cs
    • Added SessionResumptionConfig class for configuring session resumption with a Handle property.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces live session resumption and context window compression features. It adds new configuration classes (ContextWindowCompressionConfig, SlidingWindow, SessionResumptionConfig) and modifies existing classes (LiveGenerationConfig, LiveGenerativeModel, LiveSession, LiveSessionResponse) to support these functionalities. The changes involve updating method signatures, adding new fields and properties, and modifying JSON serialization/deserialization logic. While the overall implementation aligns with the new features, there are a few areas that could be improved for robustness and clarity, particularly concerning exception handling and resource management.

Comment on lines 198 to 204
catch (Exception)
{
if (clientWebSocket.State == WebSocketState.Open)
{
// Try to clean up the WebSocket, to avoid leaking connections.
await clientWebSocket.CloseAsync(WebSocketCloseStatus.EndpointUnavailable,
"Failed to send initial setup message.", CancellationToken.None);
}
// Try to clean up the WebSocket, to avoid leaking connections.
// It might not be available in scope, we rely on GC mostly here unless we catch on clientWebSocket explicitly.
// Wait, clientWebSocket is available because this is all within the lambda!
throw;
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The catch (Exception) block within the connectFactory lambda does not explicitly close the clientWebSocket if an exception occurs during the SendAsync operation. This could lead to a resource leak where an open WebSocket connection is not properly disposed of if the initial setup message fails to send.

      catch (Exception)
      {
        if (clientWebSocket.State == WebSocketState.Open)
        {
          await clientWebSocket.CloseAsync(WebSocketCloseStatus.InternalServerError, "Initial setup message failed.", CancellationToken.None);
        }
        throw;
      }

Comment on lines +200 to 203
// Try to clean up the WebSocket, to avoid leaking connections.
// It might not be available in scope, we rely on GC mostly here unless we catch on clientWebSocket explicitly.
// Wait, clientWebSocket is available because this is all within the lambda!
throw;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The comment on these lines, "// It might not be available in scope, we rely on GC mostly here unless we catch on clientWebSocket explicitly. // Wait, clientWebSocket is available because this is all within the lambda!", is misleading. clientWebSocket is clearly in scope within the lambda. This comment should be removed or updated to accurately reflect the code's behavior and the intended cleanup strategy.

        // Try to clean up the WebSocket, to avoid leaking connections.
        throw;

Comment on lines +313 to +317
catch (Exception) when (currentWebSocket != _clientWebSocket && !cancellationToken.IsCancellationRequested)
{
// The socket was closed or disposed because of session resumption, grab the new one
await Task.Delay(10, cancellationToken);
continue;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The catch (Exception) block on line 313 in ReceiveAsync is too broad. Catching a general Exception can mask specific issues that are not related to session resumption or cancellation. It would be more robust to catch specific WebSocket-related exceptions (e.g., WebSocketException, OperationCanceledException) if the intent is to handle only those scenarios.

        catch (WebSocketException) when (currentWebSocket != _clientWebSocket && !cancellationToken.IsCancellationRequested)
        {
          // The socket was closed or disposed because of session resumption, grab the new one
          await Task.Delay(10, cancellationToken);
          continue;
        }

Comment on lines +428 to +430
try {
lastConsumedClientMessageIndex = System.Convert.ToInt32(objIndex);
} catch { /* ignore */ }
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

In LiveSessionResumptionUpdate.FromJson, the try...catch block silently ignores exceptions that occur during the conversion of objIndex to an integer. This can hide malformed data in the lastConsumedClientMessageIndex field, leading to unexpected behavior or incorrect state. It's recommended to at least log the exception or handle the invalid data more explicitly.

        try {
          lastConsumedClientMessageIndex = System.Convert.ToInt32(objIndex);
        } catch (Exception e) { UnityEngine.Debug.LogError($"Failed to parse lastConsumedClientMessageIndex: {e.Message}"); }

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant