Skip to content

Commit 6db5644

Browse files
cynthiajoancynthiajoan
committed
Init setup for live session resumption
1 parent 56a084b commit 6db5644

13 files changed

+292
-15
lines changed

‎FirebaseAI/Sources/Types/Internal/Live/BidiContextWindowCompressionConfig.swift‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
import Foundation
16+
17+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
18+
@available(watchOS, unavailable)
19+
struct BidiSlidingWindow: Encodable, Sendable {
20+
let targetTokens: Int?
21+
}
22+
23+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
24+
@available(watchOS, unavailable)
25+
struct BidiContextWindowCompressionConfig: Encodable, Sendable {
26+
let triggerTokens: Int?
27+
let slidingWindow: BidiSlidingWindow?
28+
}
29+
30+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
31+
@available(watchOS, unavailable)
32+
extension SlidingWindow {
33+
var bidiSlidingWindow: BidiSlidingWindow {
34+
BidiSlidingWindow(targetTokens: targetTokens)
35+
}
36+
}
37+
38+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
39+
@available(watchOS, unavailable)
40+
extension ContextWindowCompressionConfig {
41+
var bidiContextWindowCompressionConfig: BidiContextWindowCompressionConfig {
42+
BidiContextWindowCompressionConfig(
43+
triggerTokens: triggerTokens,
44+
slidingWindow: slidingWindow?.bidiSlidingWindow
45+
)
46+
}
47+
}

‎FirebaseAI/Sources/Types/Internal/Live/BidiGenerateContentServerMessage.swift‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,9 @@ struct BidiGenerateContentServerMessage: Sendable {
3535
/// and should be cancelled.
3636
case toolCallCancellation(BidiGenerateContentToolCallCancellation)
3737

38+
/// Update with the state needed to resume the session.
39+
case sessionResumptionUpdate(BidiSessionResumptionUpdate)
40+
3841
/// Server will disconnect soon.
3942
case goAway(GoAway)
4043
}
@@ -56,6 +59,7 @@ extension BidiGenerateContentServerMessage: Decodable {
5659
case serverContent
5760
case toolCall
5861
case toolCallCancellation
62+
case sessionResumptionUpdate
5963
case goAway
6064
case usageMetadata
6165
}
@@ -83,6 +87,11 @@ extension BidiGenerateContentServerMessage: Decodable {
8387
forKey: .toolCallCancellation
8488
) {
8589
messageType = .toolCallCancellation(toolCallCancellation)
90+
} else if let sessionResumptionUpdate = try container.decodeIfPresent(
91+
BidiSessionResumptionUpdate.self,
92+
forKey: .sessionResumptionUpdate
93+
) {
94+
messageType = .sessionResumptionUpdate(sessionResumptionUpdate)
8695
} else if let goAway = try container.decodeIfPresent(GoAway.self, forKey: .goAway) {
8796
messageType = .goAway(goAway)
8897
} else {

‎FirebaseAI/Sources/Types/Internal/Live/BidiGenerateContentSetup.swift‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,20 +56,25 @@ struct BidiGenerateContentSetup: Encodable {
5656
/// turn.
5757
let outputAudioTranscription: BidiAudioTranscriptionConfig?
5858

59+
/// Configuration for the session resumption mechanism.
60+
let sessionResumption: BidiSessionResumptionConfig?
61+
5962
init(model: String,
6063
generationConfig: BidiGenerationConfig? = nil,
6164
systemInstruction: ModelContent? = nil,
6265
tools: [Tool]? = nil,
6366
toolConfig: ToolConfig? = nil,
6467
inputAudioTranscription: BidiAudioTranscriptionConfig? = nil,
65-
outputAudioTranscription: BidiAudioTranscriptionConfig? = nil) {
68+
outputAudioTranscription: BidiAudioTranscriptionConfig? = nil,
69+
sessionResumption: BidiSessionResumptionConfig? = nil) {
6670
self.model = model
6771
self.generationConfig = generationConfig
6872
self.systemInstruction = systemInstruction
6973
self.tools = tools
7074
self.toolConfig = toolConfig
7175
self.inputAudioTranscription = inputAudioTranscription
7276
self.outputAudioTranscription = outputAudioTranscription
77+
self.sessionResumption = sessionResumption
7378
}
7479
}
7580

‎FirebaseAI/Sources/Types/Internal/Live/BidiGenerationConfig.swift‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,12 +27,14 @@ struct BidiGenerationConfig: Encodable, Sendable {
2727
let frequencyPenalty: Float?
2828
let responseModalities: [ResponseModality]?
2929
let speechConfig: BidiSpeechConfig?
30+
let contextWindowCompression: BidiContextWindowCompressionConfig?
3031

3132
init(temperature: Float? = nil, topP: Float? = nil, topK: Int? = nil,
3233
candidateCount: Int? = nil, maxOutputTokens: Int? = nil,
3334
presencePenalty: Float? = nil, frequencyPenalty: Float? = nil,
3435
responseModalities: [ResponseModality]? = nil,
35-
speechConfig: BidiSpeechConfig? = nil) {
36+
speechConfig: BidiSpeechConfig? = nil,
37+
contextWindowCompression: BidiContextWindowCompressionConfig? = nil) {
3638
self.temperature = temperature
3739
self.topP = topP
3840
self.topK = topK
@@ -42,5 +44,6 @@ struct BidiGenerationConfig: Encodable, Sendable {
4244
self.frequencyPenalty = frequencyPenalty
4345
self.responseModalities = responseModalities
4446
self.speechConfig = speechConfig
47+
self.contextWindowCompression = contextWindowCompression
4548
}
4649
}

‎FirebaseAI/Sources/Types/Internal/Live/BidiSessionResumption.swift‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
import Foundation
16+
17+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
18+
@available(watchOS, unavailable)
19+
struct BidiSessionResumptionConfig: Encodable, Sendable {
20+
let handle: String?
21+
}
22+
23+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
24+
@available(watchOS, unavailable)
25+
struct BidiSessionResumptionUpdate: Decodable, Sendable {
26+
let newHandle: String?
27+
let resumable: Bool?
28+
let lastConsumedClientMessageIndex: Int?
29+
}
30+
31+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
32+
@available(watchOS, unavailable)
33+
extension SessionResumptionConfig {
34+
var bidiSessionResumptionConfig: BidiSessionResumptionConfig {
35+
BidiSessionResumptionConfig(handle: handle)
36+
}
37+
}
38+
39+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
40+
@available(watchOS, unavailable)
41+
extension LiveSessionResumptionUpdate {
42+
init(_ msg: BidiSessionResumptionUpdate) {
43+
self.init(
44+
newHandle: msg.newHandle,
45+
resumable: msg.resumable,
46+
lastConsumedClientMessageIndex: msg.lastConsumedClientMessageIndex
47+
)
48+
}
49+
}

‎FirebaseAI/Sources/Types/Internal/Live/LiveSessionService.swift‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -109,11 +109,11 @@ actor LiveSessionService {
109109
/// resuming the same session.
110110
///
111111
/// This function will yield until the websocket is ready to communicate with the client.
112-
func connect() async throws {
112+
func connect(sessionResumption: SessionResumptionConfig? = nil) async throws {
113113
close()
114114

115115
let stream = try await setupWebsocket()
116-
try await waitForSetupComplete(stream: stream)
116+
try await waitForSetupComplete(stream: stream, sessionResumption: sessionResumption)
117117
spawnMessageTasks(stream: stream)
118118
}
119119

@@ -138,10 +138,10 @@ actor LiveSessionService {
138138
/// - Server sends back `BidiGenerateContentSetupComplete` when it's ready
139139
///
140140
/// This function will yield until the setup is complete.
141-
private func waitForSetupComplete(stream: MappedStream<
142-
URLSessionWebSocketTask.Message,
143-
Data
144-
>) async throws {
141+
private func waitForSetupComplete(
142+
stream: MappedStream<URLSessionWebSocketTask.Message, Data>,
143+
sessionResumption: SessionResumptionConfig?
144+
) async throws {
145145
guard let webSocket else { return }
146146

147147
do {
@@ -152,7 +152,8 @@ actor LiveSessionService {
152152
tools: tools,
153153
toolConfig: toolConfig,
154154
inputAudioTranscription: generationConfig?.inputAudioTranscription,
155-
outputAudioTranscription: generationConfig?.outputAudioTranscription
155+
outputAudioTranscription: generationConfig?.outputAudioTranscription,
156+
sessionResumption: sessionResumption?.bidiSessionResumptionConfig
156157
)
157158
let data = try jsonEncoder.encode(BidiGenerateContentClientMessage.setup(setup))
158159
try await webSocket.send(.data(data))

‎FirebaseAI/Sources/Types/Public/Live/ContextWindowCompressionConfig.swift‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
import Foundation
16+
17+
/// Configures the sliding window context compression mechanism.
18+
///
19+
/// The context window will be truncated by keeping only a suffix of it.
20+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
21+
@available(watchOS, unavailable)
22+
public struct SlidingWindow: Sendable {
23+
/// The session reduction target, i.e., how many tokens we should keep.
24+
public let targetTokens: Int?
25+
26+
/// Creates a ``SlidingWindow`` instance.
27+
///
28+
/// - Parameter targetTokens: The target number of tokens to keep in the context window.
29+
public init(targetTokens: Int? = nil) {
30+
self.targetTokens = targetTokens
31+
}
32+
}
33+
34+
/// Enables context window compression to manage the model's context window.
35+
///
36+
/// This mechanism prevents the context from exceeding a given length.
37+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *)
38+
@available(watchOS, unavailable)
39+
public struct ContextWindowCompressionConfig: Sendable {
40+
/// The number of tokens (before running a turn) that triggers the context
41+
/// window compression.
42+
public let triggerTokens: Int?
43+
44+
/// The sliding window compression mechanism.
45+
public let slidingWindow: SlidingWindow?
46+
47+
/// Creates a ``ContextWindowCompressionConfig`` instance.
48+
///
49+
/// - Parameters:
50+
/// - triggerTokens: The number of tokens that triggers the compression mechanism.
51+
/// - slidingWindow: The sliding window compression mechanism to use.
52+
public init(triggerTokens: Int? = nil, slidingWindow: SlidingWindow? = nil) {
53+
self.triggerTokens = triggerTokens
54+
self.slidingWindow = slidingWindow
55+
}
56+
}

‎FirebaseAI/Sources/Types/Public/Live/LiveGenerationConfig.swift‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ public struct LiveGenerationConfig: Sendable {
2121
let bidiGenerationConfig: BidiGenerationConfig
2222
let inputAudioTranscription: BidiAudioTranscriptionConfig?
2323
let outputAudioTranscription: BidiAudioTranscriptionConfig?
24+
public let contextWindowCompression: ContextWindowCompressionConfig?
2425

2526
/// Creates a new ``LiveGenerationConfig`` value.
2627
///
@@ -125,7 +126,8 @@ public struct LiveGenerationConfig: Sendable {
125126
responseModalities: [ResponseModality]? = nil,
126127
speech: SpeechConfig? = nil,
127128
inputAudioTranscription: AudioTranscriptionConfig? = nil,
128-
outputAudioTranscription: AudioTranscriptionConfig? = nil) {
129+
outputAudioTranscription: AudioTranscriptionConfig? = nil,
130+
contextWindowCompression: ContextWindowCompressionConfig? = nil) {
129131
self.init(
130132
BidiGenerationConfig(
131133
temperature: temperature,
@@ -136,18 +138,22 @@ public struct LiveGenerationConfig: Sendable {
136138
presencePenalty: presencePenalty,
137139
frequencyPenalty: frequencyPenalty,
138140
responseModalities: responseModalities,
139-
speechConfig: speech?.speechConfig
141+
speechConfig: speech?.speechConfig,
142+
contextWindowCompression: contextWindowCompression?.bidiContextWindowCompressionConfig
140143
),
141144
inputAudioTranscription: inputAudioTranscription?.audioTranscriptionConfig,
142-
outputAudioTranscription: outputAudioTranscription?.audioTranscriptionConfig
145+
outputAudioTranscription: outputAudioTranscription?.audioTranscriptionConfig,
146+
contextWindowCompression: contextWindowCompression
143147
)
144148
}
145149

146150
init(_ bidiGenerationConfig: BidiGenerationConfig,
147151
inputAudioTranscription: BidiAudioTranscriptionConfig? = nil,
148-
outputAudioTranscription: BidiAudioTranscriptionConfig? = nil) {
152+
outputAudioTranscription: BidiAudioTranscriptionConfig? = nil,
153+
contextWindowCompression: ContextWindowCompressionConfig? = nil) {
149154
self.bidiGenerationConfig = bidiGenerationConfig
150155
self.inputAudioTranscription = inputAudioTranscription
151156
self.outputAudioTranscription = outputAudioTranscription
157+
self.contextWindowCompression = contextWindowCompression
152158
}
153159
}

‎FirebaseAI/Sources/Types/Public/Live/LiveGenerativeModel.swift‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,8 +53,9 @@ public final class LiveGenerativeModel {
5353

5454
/// Start a ``LiveSession`` with the server for bidirectional streaming.
5555
///
56+
/// - Parameter sessionResumption: The configuration for session resumption.
5657
/// - Returns: A new ``LiveSession`` that you can use to stream messages to and from the server.
57-
public func connect() async throws -> LiveSession {
58+
public func connect(sessionResumption: SessionResumptionConfig? = nil) async throws -> LiveSession {
5859
let service = LiveSessionService(
5960
modelResourceName: modelResourceName,
6061
generationConfig: generationConfig,
@@ -67,7 +68,7 @@ public final class LiveGenerativeModel {
6768
requestOptions: requestOptions
6869
)
6970

70-
try await service.connect()
71+
try await service.connect(sessionResumption: sessionResumption)
7172

7273
return LiveSession(service: service)
7374
}

‎FirebaseAI/Sources/Types/Public/Live/LiveServerMessage.swift‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,9 @@ public struct LiveServerMessage: Sendable {
3434
/// cancelled.
3535
case toolCallCancellation(LiveServerToolCallCancellation)
3636

37+
/// Update with the session state needed to resume the session.
38+
case sessionResumptionUpdate(LiveSessionResumptionUpdate)
39+
3740
/// Server will disconnect soon.
3841
case goingAwayNotice(LiveServerGoingAwayNotice)
3942
}
@@ -74,6 +77,8 @@ extension LiveServerMessage.Payload {
7477
self = .toolCall(LiveServerToolCall(msg))
7578
case let .toolCallCancellation(msg):
7679
self = .toolCallCancellation(LiveServerToolCallCancellation(msg))
80+
case let .sessionResumptionUpdate(msg):
81+
self = .sessionResumptionUpdate(LiveSessionResumptionUpdate(msg))
7782
case let .goAway(msg):
7883
self = .goingAwayNotice(LiveServerGoingAwayNotice(msg))
7984
}

0 commit comments

Comments
 (0)