FirebaseAI: Implement live session resumption and context window compression#1410
FirebaseAI: Implement live session resumption and context window compression#1410cynthiajoan wants to merge 1 commit intomainfrom
Conversation
…n with new configuration classes and updated connection logic.
Summary of ChangesHello, 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
🧠 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
Using Gemini Code AssistThe 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
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 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
|
There was a problem hiding this comment.
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.
| 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; | ||
| } |
There was a problem hiding this comment.
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;
}| // 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; |
There was a problem hiding this comment.
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;| 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; |
There was a problem hiding this comment.
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;
}| try { | ||
| lastConsumedClientMessageIndex = System.Convert.ToInt32(objIndex); | ||
| } catch { /* ignore */ } |
There was a problem hiding this comment.
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}"); }
…n with new configuration classes and updated connection logic.
Description
Testing
Type of Change
Place an
xthe applicable box: