diff --git a/FirebaseAI/Sources/Types/Internal/Live/BidiContextWindowCompressionConfig.swift b/FirebaseAI/Sources/Types/Internal/Live/BidiContextWindowCompressionConfig.swift new file mode 100644 index 00000000000..3c86d01d3f0 --- /dev/null +++ b/FirebaseAI/Sources/Types/Internal/Live/BidiContextWindowCompressionConfig.swift @@ -0,0 +1,47 @@ +// Copyright 2026 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 Foundation + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *) +@available(watchOS, unavailable) +struct BidiSlidingWindow: Encodable, Sendable { + let targetTokens: Int? +} + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *) +@available(watchOS, unavailable) +struct BidiContextWindowCompressionConfig: Encodable, Sendable { + let triggerTokens: Int? + let slidingWindow: BidiSlidingWindow? +} + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *) +@available(watchOS, unavailable) +extension SlidingWindow { + var bidiSlidingWindow: BidiSlidingWindow { + BidiSlidingWindow(targetTokens: targetTokens) + } +} + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *) +@available(watchOS, unavailable) +extension ContextWindowCompressionConfig { + var bidiContextWindowCompressionConfig: BidiContextWindowCompressionConfig { + BidiContextWindowCompressionConfig( + triggerTokens: triggerTokens, + slidingWindow: slidingWindow?.bidiSlidingWindow + ) + } +} diff --git a/FirebaseAI/Sources/Types/Internal/Live/BidiGenerateContentServerMessage.swift b/FirebaseAI/Sources/Types/Internal/Live/BidiGenerateContentServerMessage.swift index 8c7c628ebdb..515f6b6d889 100644 --- a/FirebaseAI/Sources/Types/Internal/Live/BidiGenerateContentServerMessage.swift +++ b/FirebaseAI/Sources/Types/Internal/Live/BidiGenerateContentServerMessage.swift @@ -35,6 +35,9 @@ struct BidiGenerateContentServerMessage: Sendable { /// and should be cancelled. case toolCallCancellation(BidiGenerateContentToolCallCancellation) + /// Update with the state needed to resume the session. + case sessionResumptionUpdate(BidiSessionResumptionUpdate) + /// Server will disconnect soon. case goAway(GoAway) } @@ -56,6 +59,7 @@ extension BidiGenerateContentServerMessage: Decodable { case serverContent case toolCall case toolCallCancellation + case sessionResumptionUpdate case goAway case usageMetadata } @@ -83,6 +87,11 @@ extension BidiGenerateContentServerMessage: Decodable { forKey: .toolCallCancellation ) { messageType = .toolCallCancellation(toolCallCancellation) + } else if let sessionResumptionUpdate = try container.decodeIfPresent( + BidiSessionResumptionUpdate.self, + forKey: .sessionResumptionUpdate + ) { + messageType = .sessionResumptionUpdate(sessionResumptionUpdate) } else if let goAway = try container.decodeIfPresent(GoAway.self, forKey: .goAway) { messageType = .goAway(goAway) } else { diff --git a/FirebaseAI/Sources/Types/Internal/Live/BidiGenerateContentSetup.swift b/FirebaseAI/Sources/Types/Internal/Live/BidiGenerateContentSetup.swift index 15dc8889a0b..1a1a72d7988 100644 --- a/FirebaseAI/Sources/Types/Internal/Live/BidiGenerateContentSetup.swift +++ b/FirebaseAI/Sources/Types/Internal/Live/BidiGenerateContentSetup.swift @@ -56,13 +56,17 @@ struct BidiGenerateContentSetup: Encodable { /// turn. let outputAudioTranscription: BidiAudioTranscriptionConfig? + /// Configuration for the session resumption mechanism. + let sessionResumption: BidiSessionResumptionConfig? + init(model: String, generationConfig: BidiGenerationConfig? = nil, systemInstruction: ModelContent? = nil, tools: [Tool]? = nil, toolConfig: ToolConfig? = nil, inputAudioTranscription: BidiAudioTranscriptionConfig? = nil, - outputAudioTranscription: BidiAudioTranscriptionConfig? = nil) { + outputAudioTranscription: BidiAudioTranscriptionConfig? = nil, + sessionResumption: BidiSessionResumptionConfig? = nil) { self.model = model self.generationConfig = generationConfig self.systemInstruction = systemInstruction @@ -70,6 +74,7 @@ struct BidiGenerateContentSetup: Encodable { self.toolConfig = toolConfig self.inputAudioTranscription = inputAudioTranscription self.outputAudioTranscription = outputAudioTranscription + self.sessionResumption = sessionResumption } } diff --git a/FirebaseAI/Sources/Types/Internal/Live/BidiGenerationConfig.swift b/FirebaseAI/Sources/Types/Internal/Live/BidiGenerationConfig.swift index a3a3e8a9f99..69c32948faf 100644 --- a/FirebaseAI/Sources/Types/Internal/Live/BidiGenerationConfig.swift +++ b/FirebaseAI/Sources/Types/Internal/Live/BidiGenerationConfig.swift @@ -27,12 +27,14 @@ struct BidiGenerationConfig: Encodable, Sendable { let frequencyPenalty: Float? let responseModalities: [ResponseModality]? let speechConfig: BidiSpeechConfig? + let contextWindowCompression: BidiContextWindowCompressionConfig? init(temperature: Float? = nil, topP: Float? = nil, topK: Int? = nil, candidateCount: Int? = nil, maxOutputTokens: Int? = nil, presencePenalty: Float? = nil, frequencyPenalty: Float? = nil, responseModalities: [ResponseModality]? = nil, - speechConfig: BidiSpeechConfig? = nil) { + speechConfig: BidiSpeechConfig? = nil, + contextWindowCompression: BidiContextWindowCompressionConfig? = nil) { self.temperature = temperature self.topP = topP self.topK = topK @@ -42,5 +44,6 @@ struct BidiGenerationConfig: Encodable, Sendable { self.frequencyPenalty = frequencyPenalty self.responseModalities = responseModalities self.speechConfig = speechConfig + self.contextWindowCompression = contextWindowCompression } } diff --git a/FirebaseAI/Sources/Types/Internal/Live/BidiSessionResumption.swift b/FirebaseAI/Sources/Types/Internal/Live/BidiSessionResumption.swift new file mode 100644 index 00000000000..308decc892a --- /dev/null +++ b/FirebaseAI/Sources/Types/Internal/Live/BidiSessionResumption.swift @@ -0,0 +1,49 @@ +// Copyright 2026 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 Foundation + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *) +@available(watchOS, unavailable) +struct BidiSessionResumptionConfig: Encodable, Sendable { + let handle: String? +} + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *) +@available(watchOS, unavailable) +struct BidiSessionResumptionUpdate: Decodable, Sendable { + let newHandle: String? + let resumable: Bool? + let lastConsumedClientMessageIndex: Int? +} + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *) +@available(watchOS, unavailable) +extension SessionResumptionConfig { + var bidiSessionResumptionConfig: BidiSessionResumptionConfig { + BidiSessionResumptionConfig(handle: handle) + } +} + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *) +@available(watchOS, unavailable) +extension LiveSessionResumptionUpdate { + init(_ msg: BidiSessionResumptionUpdate) { + self.init( + newHandle: msg.newHandle, + resumable: msg.resumable, + lastConsumedClientMessageIndex: msg.lastConsumedClientMessageIndex + ) + } +} diff --git a/FirebaseAI/Sources/Types/Internal/Live/LiveSessionService.swift b/FirebaseAI/Sources/Types/Internal/Live/LiveSessionService.swift index a2fd31b34e9..3fa496b728b 100644 --- a/FirebaseAI/Sources/Types/Internal/Live/LiveSessionService.swift +++ b/FirebaseAI/Sources/Types/Internal/Live/LiveSessionService.swift @@ -109,11 +109,11 @@ actor LiveSessionService { /// resuming the same session. /// /// This function will yield until the websocket is ready to communicate with the client. - func connect() async throws { + func connect(sessionResumption: SessionResumptionConfig? = nil) async throws { close() let stream = try await setupWebsocket() - try await waitForSetupComplete(stream: stream) + try await waitForSetupComplete(stream: stream, sessionResumption: sessionResumption) spawnMessageTasks(stream: stream) } @@ -138,10 +138,10 @@ actor LiveSessionService { /// - Server sends back `BidiGenerateContentSetupComplete` when it's ready /// /// This function will yield until the setup is complete. - private func waitForSetupComplete(stream: MappedStream< - URLSessionWebSocketTask.Message, - Data - >) async throws { + private func waitForSetupComplete( + stream: MappedStream, + sessionResumption: SessionResumptionConfig? + ) async throws { guard let webSocket else { return } do { @@ -152,7 +152,8 @@ actor LiveSessionService { tools: tools, toolConfig: toolConfig, inputAudioTranscription: generationConfig?.inputAudioTranscription, - outputAudioTranscription: generationConfig?.outputAudioTranscription + outputAudioTranscription: generationConfig?.outputAudioTranscription, + sessionResumption: sessionResumption?.bidiSessionResumptionConfig ) let data = try jsonEncoder.encode(BidiGenerateContentClientMessage.setup(setup)) try await webSocket.send(.data(data)) diff --git a/FirebaseAI/Sources/Types/Public/Live/ContextWindowCompressionConfig.swift b/FirebaseAI/Sources/Types/Public/Live/ContextWindowCompressionConfig.swift new file mode 100644 index 00000000000..4cbd42e0d83 --- /dev/null +++ b/FirebaseAI/Sources/Types/Public/Live/ContextWindowCompressionConfig.swift @@ -0,0 +1,56 @@ +// Copyright 2026 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 Foundation + +/// Configures the sliding window context compression mechanism. +/// +/// The context window will be truncated by keeping only a suffix of it. +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *) +@available(watchOS, unavailable) +public struct SlidingWindow: Sendable { + /// The session reduction target, i.e., how many tokens we should keep. + public let targetTokens: Int? + + /// Creates a ``SlidingWindow`` instance. + /// + /// - Parameter targetTokens: The target number of tokens to keep in the context window. + public init(targetTokens: Int? = nil) { + self.targetTokens = targetTokens + } +} + +/// Enables context window compression to manage the model's context window. +/// +/// This mechanism prevents the context from exceeding a given length. +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *) +@available(watchOS, unavailable) +public struct ContextWindowCompressionConfig: Sendable { + /// The number of tokens (before running a turn) that triggers the context + /// window compression. + public let triggerTokens: Int? + + /// The sliding window compression mechanism. + public let slidingWindow: SlidingWindow? + + /// Creates a ``ContextWindowCompressionConfig`` instance. + /// + /// - Parameters: + /// - triggerTokens: The number of tokens that triggers the compression mechanism. + /// - slidingWindow: The sliding window compression mechanism to use. + public init(triggerTokens: Int? = nil, slidingWindow: SlidingWindow? = nil) { + self.triggerTokens = triggerTokens + self.slidingWindow = slidingWindow + } +} diff --git a/FirebaseAI/Sources/Types/Public/Live/LiveGenerationConfig.swift b/FirebaseAI/Sources/Types/Public/Live/LiveGenerationConfig.swift index c7033567a91..8a13486641f 100644 --- a/FirebaseAI/Sources/Types/Public/Live/LiveGenerationConfig.swift +++ b/FirebaseAI/Sources/Types/Public/Live/LiveGenerationConfig.swift @@ -21,6 +21,7 @@ public struct LiveGenerationConfig: Sendable { let bidiGenerationConfig: BidiGenerationConfig let inputAudioTranscription: BidiAudioTranscriptionConfig? let outputAudioTranscription: BidiAudioTranscriptionConfig? + public let contextWindowCompression: ContextWindowCompressionConfig? /// Creates a new ``LiveGenerationConfig`` value. /// @@ -125,7 +126,8 @@ public struct LiveGenerationConfig: Sendable { responseModalities: [ResponseModality]? = nil, speech: SpeechConfig? = nil, inputAudioTranscription: AudioTranscriptionConfig? = nil, - outputAudioTranscription: AudioTranscriptionConfig? = nil) { + outputAudioTranscription: AudioTranscriptionConfig? = nil, + contextWindowCompression: ContextWindowCompressionConfig? = nil) { self.init( BidiGenerationConfig( temperature: temperature, @@ -136,18 +138,22 @@ public struct LiveGenerationConfig: Sendable { presencePenalty: presencePenalty, frequencyPenalty: frequencyPenalty, responseModalities: responseModalities, - speechConfig: speech?.speechConfig + speechConfig: speech?.speechConfig, + contextWindowCompression: contextWindowCompression?.bidiContextWindowCompressionConfig ), inputAudioTranscription: inputAudioTranscription?.audioTranscriptionConfig, - outputAudioTranscription: outputAudioTranscription?.audioTranscriptionConfig + outputAudioTranscription: outputAudioTranscription?.audioTranscriptionConfig, + contextWindowCompression: contextWindowCompression ) } init(_ bidiGenerationConfig: BidiGenerationConfig, inputAudioTranscription: BidiAudioTranscriptionConfig? = nil, - outputAudioTranscription: BidiAudioTranscriptionConfig? = nil) { + outputAudioTranscription: BidiAudioTranscriptionConfig? = nil, + contextWindowCompression: ContextWindowCompressionConfig? = nil) { self.bidiGenerationConfig = bidiGenerationConfig self.inputAudioTranscription = inputAudioTranscription self.outputAudioTranscription = outputAudioTranscription + self.contextWindowCompression = contextWindowCompression } } diff --git a/FirebaseAI/Sources/Types/Public/Live/LiveGenerativeModel.swift b/FirebaseAI/Sources/Types/Public/Live/LiveGenerativeModel.swift index 3a8236cb1d5..a062620582f 100644 --- a/FirebaseAI/Sources/Types/Public/Live/LiveGenerativeModel.swift +++ b/FirebaseAI/Sources/Types/Public/Live/LiveGenerativeModel.swift @@ -53,8 +53,9 @@ public final class LiveGenerativeModel { /// Start a ``LiveSession`` with the server for bidirectional streaming. /// + /// - Parameter sessionResumption: The configuration for session resumption. /// - Returns: A new ``LiveSession`` that you can use to stream messages to and from the server. - public func connect() async throws -> LiveSession { + public func connect(sessionResumption: SessionResumptionConfig? = nil) async throws -> LiveSession { let service = LiveSessionService( modelResourceName: modelResourceName, generationConfig: generationConfig, @@ -67,7 +68,7 @@ public final class LiveGenerativeModel { requestOptions: requestOptions ) - try await service.connect() + try await service.connect(sessionResumption: sessionResumption) return LiveSession(service: service) } diff --git a/FirebaseAI/Sources/Types/Public/Live/LiveServerMessage.swift b/FirebaseAI/Sources/Types/Public/Live/LiveServerMessage.swift index af6caca90c1..97f93728f50 100644 --- a/FirebaseAI/Sources/Types/Public/Live/LiveServerMessage.swift +++ b/FirebaseAI/Sources/Types/Public/Live/LiveServerMessage.swift @@ -34,6 +34,9 @@ public struct LiveServerMessage: Sendable { /// cancelled. case toolCallCancellation(LiveServerToolCallCancellation) + /// Update with the session state needed to resume the session. + case sessionResumptionUpdate(LiveSessionResumptionUpdate) + /// Server will disconnect soon. case goingAwayNotice(LiveServerGoingAwayNotice) } @@ -74,6 +77,8 @@ extension LiveServerMessage.Payload { self = .toolCall(LiveServerToolCall(msg)) case let .toolCallCancellation(msg): self = .toolCallCancellation(LiveServerToolCallCancellation(msg)) + case let .sessionResumptionUpdate(msg): + self = .sessionResumptionUpdate(LiveSessionResumptionUpdate(msg)) case let .goAway(msg): self = .goingAwayNotice(LiveServerGoingAwayNotice(msg)) } diff --git a/FirebaseAI/Sources/Types/Public/Live/LiveSession.swift b/FirebaseAI/Sources/Types/Public/Live/LiveSession.swift index a7ece57bdfc..1e81a6b46b0 100644 --- a/FirebaseAI/Sources/Types/Public/Live/LiveSession.swift +++ b/FirebaseAI/Sources/Types/Public/Live/LiveSession.swift @@ -147,4 +147,18 @@ public final class LiveSession: Sendable { } // TODO: b(445716402) Add a start method when we support session resumption + + + /// Resumes an existing live session with the server. + /// + /// This closes the current WebSocket connection and establishes a new one using + /// the same configuration (URI, headers, model, system instruction, tools, etc.) + /// as the original session. + /// + /// - Parameter sessionResumption: The configuration for session resumption, + /// such as the handle to the previous session state to restore. + public func resumeSession(sessionResumption: SessionResumptionConfig? = nil) async throws { + try await service.connect(sessionResumption: sessionResumption) + } } + diff --git a/FirebaseAI/Sources/Types/Public/Live/LiveSessionResumptionUpdate.swift b/FirebaseAI/Sources/Types/Public/Live/LiveSessionResumptionUpdate.swift new file mode 100644 index 00000000000..d37a694fa0e --- /dev/null +++ b/FirebaseAI/Sources/Types/Public/Live/LiveSessionResumptionUpdate.swift @@ -0,0 +1,46 @@ +// Copyright 2026 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 Foundation + +/// An update of the session resumption state. +/// +/// This message is only sent if ``SessionResumptionConfig`` was set in the +/// session setup. +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *) +@available(watchOS, unavailable) +public struct LiveSessionResumptionUpdate: Sendable { + /// The new handle that represents the state that can be resumed. Empty if + /// `resumable` is false. + public let newHandle: String? + + /// Indicates if the session can be resumed at this point. + public let resumable: Bool? + + /// The index of the last client message that is included in the state + /// represented by this update. + public let lastConsumedClientMessageIndex: Int? + + /// Creates a ``LiveSessionResumptionUpdate`` instance. + /// + /// - Parameters: + /// - newHandle: The new handle that represents the state that can be resumed. + /// - resumable: Indicates if the session can be resumed at this point. + /// - lastConsumedClientMessageIndex: The index of the last client message that is included in the state. + public init(newHandle: String? = nil, resumable: Bool? = nil, lastConsumedClientMessageIndex: Int? = nil) { + self.newHandle = newHandle + self.resumable = resumable + self.lastConsumedClientMessageIndex = lastConsumedClientMessageIndex + } +} diff --git a/FirebaseAI/Sources/Types/Public/Live/SessionResumptionConfig.swift b/FirebaseAI/Sources/Types/Public/Live/SessionResumptionConfig.swift new file mode 100644 index 00000000000..8ec2201f93b --- /dev/null +++ b/FirebaseAI/Sources/Types/Public/Live/SessionResumptionConfig.swift @@ -0,0 +1,35 @@ +// Copyright 2026 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 Foundation + +/// Configuration for the session resumption mechanism. +/// +/// When included in the session setup, the server will send +/// ``LiveSessionResumptionUpdate`` messages in the response stream. +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, *) +@available(watchOS, unavailable) +public struct SessionResumptionConfig: Sendable { + /// The session resumption handle of the previous session to restore. + /// + /// If not present, a new session will be started. + public let handle: String? + + /// Creates a ``SessionResumptionConfig`` instance. + /// + /// - Parameter handle: The session resumption handle of the previous session to restore. + public init(handle: String? = nil) { + self.handle = handle + } +}