From fb6ca664620c1bb650a9f8a67c7c96b7593a52a9 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Fri, 3 Oct 2025 16:27:57 -0700 Subject: [PATCH 01/24] [AI] Server Prompt Templates --- .../Sources/BaseTemplateAPIClientModel.swift | 26 +++++++ FirebaseAI/Sources/FirebaseAI.swift | 20 +++++ FirebaseAI/Sources/TemplateChatSession.swift | 44 +++++++++++ .../Sources/TemplateGenerativeModel.swift | 73 +++++++++++++++++++ FirebaseAI/Sources/TemplateImagenModel.swift | 41 +++++++++++ .../Tests/Unit/TemplateChatSessionTests.swift | 42 +++++++++++ .../Tests/Unit/TemplateImagenModelTests.swift | 39 ++++++++++ 7 files changed, 285 insertions(+) create mode 100644 FirebaseAI/Sources/BaseTemplateAPIClientModel.swift create mode 100644 FirebaseAI/Sources/TemplateChatSession.swift create mode 100644 FirebaseAI/Sources/TemplateGenerativeModel.swift create mode 100644 FirebaseAI/Sources/TemplateImagenModel.swift create mode 100644 FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift create mode 100644 FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift diff --git a/FirebaseAI/Sources/BaseTemplateAPIClientModel.swift b/FirebaseAI/Sources/BaseTemplateAPIClientModel.swift new file mode 100644 index 00000000000..0335f28ccbd --- /dev/null +++ b/FirebaseAI/Sources/BaseTemplateAPIClientModel.swift @@ -0,0 +1,26 @@ + +// Copyright 2024 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 + +/// A base class for template API client models. +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +public class BaseTemplateAPIClientModel { + let generativeAIService: GenerativeAIService + + init(generativeAIService: GenerativeAIService) { + self.generativeAIService = generativeAIService + } +} diff --git a/FirebaseAI/Sources/FirebaseAI.swift b/FirebaseAI/Sources/FirebaseAI.swift index f9ff5ea0424..6be9f4d623c 100644 --- a/FirebaseAI/Sources/FirebaseAI.swift +++ b/FirebaseAI/Sources/FirebaseAI.swift @@ -136,6 +136,26 @@ public final class FirebaseAI: Sendable { ) } + /// Initializes a new `TemplateGenerativeModel`. + /// + /// - Returns: A new `TemplateGenerativeModel` instance. + public func templateGenerativeModel() -> TemplateGenerativeModel { + return TemplateGenerativeModel(generativeAIService: GenerativeAIService( + firebaseInfo: firebaseInfo, + urlSession: GenAIURLSession.default + )) + } + + /// Initializes a new `TemplateImagenModel`. + /// + /// - Returns: A new `TemplateImagenModel` instance. + public func templateImagenModel() -> TemplateImagenModel { + return TemplateImagenModel(generativeAIService: GenerativeAIService( + firebaseInfo: firebaseInfo, + urlSession: GenAIURLSession.default + )) + } + /// **[Public Preview]** Initializes a ``LiveGenerativeModel`` with the given parameters. /// /// > Warning: Using the Firebase AI Logic SDKs with the Gemini Live API is in Public diff --git a/FirebaseAI/Sources/TemplateChatSession.swift b/FirebaseAI/Sources/TemplateChatSession.swift new file mode 100644 index 00000000000..176d951a9fd --- /dev/null +++ b/FirebaseAI/Sources/TemplateChatSession.swift @@ -0,0 +1,44 @@ + +// Copyright 2024 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 + +/// A chat session that allows for conversation with a model. +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +public class TemplateChatSession { + private let templateGenerateContent: ([ModelContent], String, [String: Any]) async throws + -> GenerateContentResponse + private let template: String + public var history: [ModelContent] + + init(templateGenerateContent: @escaping (([ModelContent], String, [String: Any]) async throws + -> GenerateContentResponse), + template: String, history: [ModelContent]) { + self.templateGenerateContent = templateGenerateContent + self.template = template + self.history = history + } + + /// Sends a message to the model and returns the response. + public func sendMessage(_ message: any PartsRepresentable, + variables: [String: Any]) async throws -> GenerateContentResponse { + let response = try await templateGenerateContent(history, template, variables) + history.append(ModelContent(role: "user", parts: message.partsValue)) + if let modelResponse = response.candidates.first { + history.append(modelResponse.content) + } + return response + } +} diff --git a/FirebaseAI/Sources/TemplateGenerativeModel.swift b/FirebaseAI/Sources/TemplateGenerativeModel.swift new file mode 100644 index 00000000000..9c44aca5d86 --- /dev/null +++ b/FirebaseAI/Sources/TemplateGenerativeModel.swift @@ -0,0 +1,73 @@ + +// Copyright 2024 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 + +/// A type that represents a remote multimodal model (like Gemini), with the ability to generate +/// content based on various input types. +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +public final class TemplateGenerativeModel: BaseTemplateAPIClientModel { + /// Generates content from a prompt template and variables. + /// + /// - Parameters: + /// - template: The prompt template to use. + /// - variables: A dictionary of variables to substitute into the template. + /// - Returns: The content generated by the model. + /// - Throws: A ``GenerateContentError`` if the request failed. + public func templateGenerateContent(template: String, + variables: [String: Any]) async throws + -> GenerateContentResponse { + return try await templateGenerateContentWithHistory( + history: [], + template: template, + variables: variables + ) + } + + /// Generates content from a prompt template, variables, and history. + /// + /// - Parameters: + /// - history: The conversation history to use. + /// - template: The prompt template to use. + /// - variables: A dictionary of variables to substitute into the template. + /// - Returns: The content generated by the model. + /// - Throws: A ``GenerateContentError`` if the request failed. + func templateGenerateContentWithHistory(history: [ModelContent], template: String, + variables: [String: Any]) async throws + -> GenerateContentResponse { + let candidate = Candidate( + content: ModelContent(role: "model", parts: ["response"]), + safetyRatings: [], + finishReason: .stop, + citationMetadata: nil + ) + return GenerateContentResponse(candidates: [candidate], promptFeedback: nil) + } + + /// Creates a new chat conversation using this model with the provided history and template. + /// + /// - Parameters: + /// - template: The prompt template to use. + /// - history: The conversation history to use. + /// - Returns: A new ``TemplateChatSession`` instance. + public func startTemplateChat(template: String, + history: [ModelContent] = []) -> TemplateChatSession { + return TemplateChatSession( + templateGenerateContent: templateGenerateContentWithHistory, + template: template, + history: history + ) + } +} diff --git a/FirebaseAI/Sources/TemplateImagenModel.swift b/FirebaseAI/Sources/TemplateImagenModel.swift new file mode 100644 index 00000000000..5e15ff77d54 --- /dev/null +++ b/FirebaseAI/Sources/TemplateImagenModel.swift @@ -0,0 +1,41 @@ + +// Copyright 2024 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 + +/// A type that represents a remote image generation model (like Imagen), with the ability to +/// generate +/// images based on various input types. +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +public final class TemplateImagenModel: BaseTemplateAPIClientModel { + /// Generates images from a prompt template and variables. + /// + /// - Parameters: + /// - template: The prompt template to use. + /// - variables: A dictionary of variables to substitute into the template. + /// - Returns: The images generated by the model. + /// - Throws: An error if the request failed. + public func templateImagenGenerateImages(template: String, + variables: [String: Any]) async throws + -> GenerateImagesResponse { + // Not implemented + return GenerateImagesResponse(images: []) + } +} + +// A placeholder for the response from an image generation request. +public struct GenerateImagesResponse { + public let images: [Data] +} diff --git a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift new file mode 100644 index 00000000000..27c9db16708 --- /dev/null +++ b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift @@ -0,0 +1,42 @@ + +// Copyright 2024 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. + +@testable import FirebaseAI +import FirebaseCore +import XCTest + +final class TemplateChatSessionTests: XCTestCase { + var model: TemplateGenerativeModel! + + override func setUp() { + super.setUp() + let firebaseInfo = GenerativeModelTestUtil.testFirebaseInfo() + let generativeAIService = GenerativeAIService( + firebaseInfo: firebaseInfo, + urlSession: GenAIURLSession.default + ) + model = TemplateGenerativeModel(generativeAIService: generativeAIService) + } + + func testSendMessage() async throws { + let chat = model.startTemplateChat(template: "test-template") + let response = try await chat.sendMessage("Hello", variables: ["name": "test"]) + XCTAssertEqual(chat.history.count, 2) + XCTAssertEqual(chat.history[0].role, "user") + XCTAssertEqual((chat.history[0].parts.first as? TextPart)?.text, "Hello") + XCTAssertEqual(chat.history[1].role, "model") + XCTAssertEqual(response.candidates.count, 1) + } +} diff --git a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift new file mode 100644 index 00000000000..8fbf7d5a2a6 --- /dev/null +++ b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift @@ -0,0 +1,39 @@ + +// Copyright 2024 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 of 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. + +@testable import FirebaseAI +import XCTest + +final class TemplateImagenModelTests: XCTestCase { + var model: TemplateImagenModel! + + override func setUp() { + super.setUp() + let firebaseInfo = GenerativeModelTestUtil.testFirebaseInfo() + let generativeAIService = GenerativeAIService( + firebaseInfo: firebaseInfo, + urlSession: GenAIURLSession.default + ) + model = TemplateImagenModel(generativeAIService: generativeAIService) + } + + func testTemplateImagenGenerateImages() async throws { + let response = try await model.templateImagenGenerateImages( + template: "test-template", + variables: ["prompt": "a cat picture"] + ) + XCTAssertEqual(response.images.count, 0) + } +} From 114478ad5cbf7190a4e9e87e3fe6eed90f7a1c92 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Fri, 3 Oct 2025 16:32:56 -0700 Subject: [PATCH 02/24] copyrights --- FirebaseAI/Sources/BaseTemplateAPIClientModel.swift | 2 +- FirebaseAI/Sources/TemplateGenerativeModel.swift | 2 +- FirebaseAI/Sources/TemplateImagenModel.swift | 2 +- FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift | 2 +- FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/FirebaseAI/Sources/BaseTemplateAPIClientModel.swift b/FirebaseAI/Sources/BaseTemplateAPIClientModel.swift index 0335f28ccbd..1f2deab1ff1 100644 --- a/FirebaseAI/Sources/BaseTemplateAPIClientModel.swift +++ b/FirebaseAI/Sources/BaseTemplateAPIClientModel.swift @@ -1,5 +1,5 @@ -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. diff --git a/FirebaseAI/Sources/TemplateGenerativeModel.swift b/FirebaseAI/Sources/TemplateGenerativeModel.swift index 9c44aca5d86..0b80a30641a 100644 --- a/FirebaseAI/Sources/TemplateGenerativeModel.swift +++ b/FirebaseAI/Sources/TemplateGenerativeModel.swift @@ -1,5 +1,5 @@ -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. diff --git a/FirebaseAI/Sources/TemplateImagenModel.swift b/FirebaseAI/Sources/TemplateImagenModel.swift index 5e15ff77d54..eb5ca008eb9 100644 --- a/FirebaseAI/Sources/TemplateImagenModel.swift +++ b/FirebaseAI/Sources/TemplateImagenModel.swift @@ -1,5 +1,5 @@ -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. diff --git a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift index 27c9db16708..1547e852e9d 100644 --- a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift @@ -1,5 +1,5 @@ -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. diff --git a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift index 8fbf7d5a2a6..8ee33da3f2c 100644 --- a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift @@ -1,5 +1,5 @@ -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. From cbdd6355307b201b35b5752256abce6fb5919b60 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Tue, 7 Oct 2025 15:58:58 -0700 Subject: [PATCH 03/24] Checkpoint after trying to add imagen implementation --- FirebaseAI/Sources/APIMethod.swift | 22 ++++++++ FirebaseAI/Sources/AnyCodable.swift | 47 +++++++++++++++++ .../Sources/BaseTemplateAPIClientModel.swift | 4 +- FirebaseAI/Sources/FirebaseAI.swift | 18 ++++--- .../Sources/GenerateContentRequest.swift | 9 ---- .../Sources/GenerateImagesRequest.swift | 51 +++++++++++++++++++ FirebaseAI/Sources/ImageAPIMethod.swift | 20 ++++++++ .../Sources/TemplateGenerativeModel.swift | 4 ++ FirebaseAI/Sources/TemplateImagenModel.swift | 34 +++++++++++-- .../Tests/Unit/TemplateChatSessionTests.swift | 3 +- .../Tests/Unit/TemplateImagenModelTests.swift | 19 +++++-- .../GenerativeModelTestUtil.swift | 7 ++- 12 files changed, 210 insertions(+), 28 deletions(-) create mode 100644 FirebaseAI/Sources/APIMethod.swift create mode 100644 FirebaseAI/Sources/AnyCodable.swift create mode 100644 FirebaseAI/Sources/GenerateImagesRequest.swift create mode 100644 FirebaseAI/Sources/ImageAPIMethod.swift diff --git a/FirebaseAI/Sources/APIMethod.swift b/FirebaseAI/Sources/APIMethod.swift new file mode 100644 index 00000000000..9ae05e88429 --- /dev/null +++ b/FirebaseAI/Sources/APIMethod.swift @@ -0,0 +1,22 @@ + +// Copyright 2025 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 + +enum APIMethod: String { + case generateContent + case streamGenerateContent + case countTokens +} diff --git a/FirebaseAI/Sources/AnyCodable.swift b/FirebaseAI/Sources/AnyCodable.swift new file mode 100644 index 00000000000..43110196c4d --- /dev/null +++ b/FirebaseAI/Sources/AnyCodable.swift @@ -0,0 +1,47 @@ + +// Copyright 2025 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 + +struct AnyCodable: Encodable { + let value: Any + + init(_ value: Any) { + self.value = value + } + + func encode(to encoder: Encoder) throws { + var container = encoder.singleValueContainer() + switch value { + case let value as String: + try container.encode(value) + case let value as Int: + try container.encode(value) + case let value as Double: + try container.encode(value) + case let value as Bool: + try container.encode(value) + case let value as [Any]: + try container.encode(value.map { AnyCodable($0) }) + case let value as [String: Any]: + try container.encode(value.mapValues { AnyCodable($0) }) + default: + throw EncodingError.invalidValue( + value, + EncodingError.Context(codingPath: [], debugDescription: "Invalid value") + ) + } + } +} diff --git a/FirebaseAI/Sources/BaseTemplateAPIClientModel.swift b/FirebaseAI/Sources/BaseTemplateAPIClientModel.swift index 1f2deab1ff1..6363310ce5b 100644 --- a/FirebaseAI/Sources/BaseTemplateAPIClientModel.swift +++ b/FirebaseAI/Sources/BaseTemplateAPIClientModel.swift @@ -19,8 +19,10 @@ import Foundation @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) public class BaseTemplateAPIClientModel { let generativeAIService: GenerativeAIService + let apiConfig: APIConfig - init(generativeAIService: GenerativeAIService) { + init(generativeAIService: GenerativeAIService, apiConfig: APIConfig) { self.generativeAIService = generativeAIService + self.apiConfig = apiConfig } } diff --git a/FirebaseAI/Sources/FirebaseAI.swift b/FirebaseAI/Sources/FirebaseAI.swift index 6be9f4d623c..2c041a876f1 100644 --- a/FirebaseAI/Sources/FirebaseAI.swift +++ b/FirebaseAI/Sources/FirebaseAI.swift @@ -140,20 +140,22 @@ public final class FirebaseAI: Sendable { /// /// - Returns: A new `TemplateGenerativeModel` instance. public func templateGenerativeModel() -> TemplateGenerativeModel { - return TemplateGenerativeModel(generativeAIService: GenerativeAIService( - firebaseInfo: firebaseInfo, - urlSession: GenAIURLSession.default - )) + return TemplateGenerativeModel( + generativeAIService: GenerativeAIService(firebaseInfo: firebaseInfo, + urlSession: GenAIURLSession.default), + apiConfig: apiConfig + ) } /// Initializes a new `TemplateImagenModel`. /// /// - Returns: A new `TemplateImagenModel` instance. public func templateImagenModel() -> TemplateImagenModel { - return TemplateImagenModel(generativeAIService: GenerativeAIService( - firebaseInfo: firebaseInfo, - urlSession: GenAIURLSession.default - )) + return TemplateImagenModel( + generativeAIService: GenerativeAIService(firebaseInfo: firebaseInfo, + urlSession: GenAIURLSession.default), + apiConfig: apiConfig + ) } /// **[Public Preview]** Initializes a ``LiveGenerativeModel`` with the given parameters. diff --git a/FirebaseAI/Sources/GenerateContentRequest.swift b/FirebaseAI/Sources/GenerateContentRequest.swift index 21acd502a75..ddf8d7d28cb 100644 --- a/FirebaseAI/Sources/GenerateContentRequest.swift +++ b/FirebaseAI/Sources/GenerateContentRequest.swift @@ -60,15 +60,6 @@ extension GenerateContentRequest: Encodable { } } -@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) -extension GenerateContentRequest { - enum APIMethod: String { - case generateContent - case streamGenerateContent - case countTokens - } -} - @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) extension GenerateContentRequest: GenerativeAIRequest { typealias Response = GenerateContentResponse diff --git a/FirebaseAI/Sources/GenerateImagesRequest.swift b/FirebaseAI/Sources/GenerateImagesRequest.swift new file mode 100644 index 00000000000..53f0c78df6a --- /dev/null +++ b/FirebaseAI/Sources/GenerateImagesRequest.swift @@ -0,0 +1,51 @@ + +// Copyright 2025 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, watchOS 8.0, *) +public class GenerateImagesRequest: @unchecked Sendable, GenerativeAIRequest { + public typealias Response = GenerateImagesResponse + + public let url: URL + public let options: RequestOptions + + let apiConfig: APIConfig + + let template: String + let variables: [String: AnyCodable] + + init(template: String, variables: [String: AnyCodable], apiConfig: APIConfig, + options: RequestOptions) { + let modelURL = + "\(apiConfig.service.endpoint.rawValue)/\(apiConfig.version.rawValue)/\(template)" + url = URL(string: "\(modelURL):\(ImageAPIMethod.generateImages.rawValue)")! + self.apiConfig = apiConfig + self.options = options + self.template = template + self.variables = variables + } + + enum CodingKeys: String, CodingKey { + case template + case variables + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(template, forKey: .template) + try container.encode(variables, forKey: .variables) + } +} diff --git a/FirebaseAI/Sources/ImageAPIMethod.swift b/FirebaseAI/Sources/ImageAPIMethod.swift new file mode 100644 index 00000000000..84afa50a742 --- /dev/null +++ b/FirebaseAI/Sources/ImageAPIMethod.swift @@ -0,0 +1,20 @@ + +// Copyright 2025 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 not use this file except in compliance with the License. +// +// 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 + +enum ImageAPIMethod: String { + case generateImages = "images:generate" +} diff --git a/FirebaseAI/Sources/TemplateGenerativeModel.swift b/FirebaseAI/Sources/TemplateGenerativeModel.swift index 0b80a30641a..90df6d2d0f8 100644 --- a/FirebaseAI/Sources/TemplateGenerativeModel.swift +++ b/FirebaseAI/Sources/TemplateGenerativeModel.swift @@ -19,6 +19,10 @@ import Foundation /// content based on various input types. @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) public final class TemplateGenerativeModel: BaseTemplateAPIClientModel { + override init(generativeAIService: GenerativeAIService, apiConfig: APIConfig) { + super.init(generativeAIService: generativeAIService, apiConfig: apiConfig) + } + /// Generates content from a prompt template and variables. /// /// - Parameters: diff --git a/FirebaseAI/Sources/TemplateImagenModel.swift b/FirebaseAI/Sources/TemplateImagenModel.swift index eb5ca008eb9..ca93f9a74d2 100644 --- a/FirebaseAI/Sources/TemplateImagenModel.swift +++ b/FirebaseAI/Sources/TemplateImagenModel.swift @@ -28,14 +28,40 @@ public final class TemplateImagenModel: BaseTemplateAPIClientModel { /// - Returns: The images generated by the model. /// - Throws: An error if the request failed. public func templateImagenGenerateImages(template: String, - variables: [String: Any]) async throws + variables: [String: Any], + options: RequestOptions = RequestOptions()) async throws -> GenerateImagesResponse { - // Not implemented - return GenerateImagesResponse(images: []) + let request = GenerateImagesRequest( + template: template, + variables: variables.mapValues { AnyCodable($0) }, + apiConfig: apiConfig, + options: options + ) + let response: GenerateImagesResponse = try await generativeAIService + .loadRequest(request: request) + return response } } // A placeholder for the response from an image generation request. -public struct GenerateImagesResponse { +public struct GenerateImagesResponse: Decodable, @unchecked Sendable { public let images: [Data] + + enum CodingKeys: String, CodingKey { + case predictions + } + + struct Prediction: Decodable { + let images: [String] + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + let predictions = try container.decode([Prediction].self, forKey: .predictions) + guard let firstPrediction = predictions.first else { + images = [] + return + } + images = firstPrediction.images.compactMap { Data(base64Encoded: $0) } + } } diff --git a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift index 1547e852e9d..d3cdff52168 100644 --- a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift @@ -27,7 +27,8 @@ final class TemplateChatSessionTests: XCTestCase { firebaseInfo: firebaseInfo, urlSession: GenAIURLSession.default ) - model = TemplateGenerativeModel(generativeAIService: generativeAIService) + let apiConfig = APIConfig(service: .googleAI(endpoint: .firebaseProxyProd), version: .v1beta) + model = TemplateGenerativeModel(generativeAIService: generativeAIService, apiConfig: apiConfig) } func testSendMessage() async throws { diff --git a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift index 8ee33da3f2c..2abd1f3d685 100644 --- a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift @@ -17,23 +17,36 @@ import XCTest final class TemplateImagenModelTests: XCTestCase { + var urlSession: URLSession! var model: TemplateImagenModel! override func setUp() { super.setUp() + let configuration = URLSessionConfiguration.default + configuration.protocolClasses = [MockURLProtocol.self] + urlSession = URLSession(configuration: configuration) let firebaseInfo = GenerativeModelTestUtil.testFirebaseInfo() let generativeAIService = GenerativeAIService( firebaseInfo: firebaseInfo, - urlSession: GenAIURLSession.default + urlSession: urlSession ) - model = TemplateImagenModel(generativeAIService: generativeAIService) + let apiConfig = APIConfig(service: .googleAI(endpoint: .firebaseProxyProd), version: .v1beta) + model = TemplateImagenModel(generativeAIService: generativeAIService, apiConfig: apiConfig) } func testTemplateImagenGenerateImages() async throws { + MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( + forResource: "unary-success-image-response", + withExtension: "json", + subdirectory: "vertexai-sdk-test-data/mock-responses", + isImageRequest: true + ) + let response = try await model.templateImagenGenerateImages( template: "test-template", variables: ["prompt": "a cat picture"] ) - XCTAssertEqual(response.images.count, 0) + XCTAssertEqual(response.images.count, 1) + XCTAssertEqual(response.images.first, Data(base64Encoded: "aW1hZ2UgZGF0YQ==")) } } diff --git a/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift b/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift index ee4f47bc5b0..aa2a67c2d34 100644 --- a/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift +++ b/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift @@ -30,7 +30,8 @@ enum GenerativeModelTestUtil { timeout: TimeInterval = RequestOptions().timeout, appCheckToken: String? = nil, authToken: String? = nil, - dataCollection: Bool = true) throws -> ((URLRequest) throws -> ( + dataCollection: Bool = true, + isImageRequest: Bool = false) throws -> ((URLRequest) throws -> ( URLResponse, AsyncLineSequence? )) { @@ -45,7 +46,9 @@ enum GenerativeModelTestUtil { ) return { request in let requestURL = try XCTUnwrap(request.url) - XCTAssertEqual(requestURL.path.occurrenceCount(of: "models/"), 1) + if !isImageRequest { + XCTAssertEqual(requestURL.path.occurrenceCount(of: "models/"), 1) + } XCTAssertEqual(request.timeoutInterval, timeout) let apiClientTags = try XCTUnwrap(request.value(forHTTPHeaderField: "x-goog-api-client")) .components(separatedBy: " ") From 91e4461025d182d28fe4f403fa622ae04e4ab58e Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Thu, 9 Oct 2025 14:23:25 -0700 Subject: [PATCH 04/24] end to end testGenerateContentWithText passes --- .../Sources/GenerateImagesRequest.swift | 4 +- FirebaseAI/Sources/GenerativeAIService.swift | 2 +- FirebaseAI/Sources/TemplateChatSession.swift | 15 ++- .../TemplateGenerateContentRequest.swift | 56 +++++++++ .../Sources/TemplateGenerativeModel.swift | 38 +++--- FirebaseAI/Sources/TemplateImagenModel.swift | 9 +- ...nyCodable.swift => TemplateVariable.swift} | 46 +++++--- ...ServerPromptTemplateIntegrationTests.swift | 108 ++++++++++++++++++ .../Tests/Unit/TemplateChatSessionTests.swift | 2 +- .../Tests/Unit/TemplateImagenModelTests.swift | 6 +- 10 files changed, 240 insertions(+), 46 deletions(-) create mode 100644 FirebaseAI/Sources/TemplateGenerateContentRequest.swift rename FirebaseAI/Sources/{AnyCodable.swift => TemplateVariable.swift} (61%) create mode 100644 FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift diff --git a/FirebaseAI/Sources/GenerateImagesRequest.swift b/FirebaseAI/Sources/GenerateImagesRequest.swift index 53f0c78df6a..81fe0f5dde8 100644 --- a/FirebaseAI/Sources/GenerateImagesRequest.swift +++ b/FirebaseAI/Sources/GenerateImagesRequest.swift @@ -25,9 +25,9 @@ public class GenerateImagesRequest: @unchecked Sendable, GenerativeAIRequest { let apiConfig: APIConfig let template: String - let variables: [String: AnyCodable] + let variables: [String: TemplateVariable] - init(template: String, variables: [String: AnyCodable], apiConfig: APIConfig, + init(template: String, variables: [String: TemplateVariable], apiConfig: APIConfig, options: RequestOptions) { let modelURL = "\(apiConfig.service.endpoint.rawValue)/\(apiConfig.version.rawValue)/\(template)" diff --git a/FirebaseAI/Sources/GenerativeAIService.swift b/FirebaseAI/Sources/GenerativeAIService.swift index a17364f8cb6..8eeef564e4c 100644 --- a/FirebaseAI/Sources/GenerativeAIService.swift +++ b/FirebaseAI/Sources/GenerativeAIService.swift @@ -26,7 +26,7 @@ struct GenerativeAIService { /// The Firebase SDK version in the format `fire/`. static let firebaseVersionTag = "fire/\(FirebaseVersion())" - private let firebaseInfo: FirebaseInfo + let firebaseInfo: FirebaseInfo private let urlSession: URLSession diff --git a/FirebaseAI/Sources/TemplateChatSession.swift b/FirebaseAI/Sources/TemplateChatSession.swift index 176d951a9fd..cbb4da9cb2d 100644 --- a/FirebaseAI/Sources/TemplateChatSession.swift +++ b/FirebaseAI/Sources/TemplateChatSession.swift @@ -18,23 +18,26 @@ import Foundation /// A chat session that allows for conversation with a model. @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) public class TemplateChatSession { - private let templateGenerateContent: ([ModelContent], String, [String: Any]) async throws + private let generateContent: ([ModelContent], String, [String: Any], RequestOptions) async throws -> GenerateContentResponse private let template: String public var history: [ModelContent] - init(templateGenerateContent: @escaping (([ModelContent], String, [String: Any]) async throws + init(generateContent: @escaping (([ModelContent], String, [String: Any], + RequestOptions) async throws -> GenerateContentResponse), - template: String, history: [ModelContent]) { - self.templateGenerateContent = templateGenerateContent + template: String, history: [ModelContent]) { + self.generateContent = generateContent self.template = template self.history = history } /// Sends a message to the model and returns the response. public func sendMessage(_ message: any PartsRepresentable, - variables: [String: Any]) async throws -> GenerateContentResponse { - let response = try await templateGenerateContent(history, template, variables) + variables: [String: Any], + options: RequestOptions = RequestOptions()) async throws + -> GenerateContentResponse { + let response = try await generateContent(history, template, variables, options) history.append(ModelContent(role: "user", parts: message.partsValue)) if let modelResponse = response.candidates.first { history.append(modelResponse.content) diff --git a/FirebaseAI/Sources/TemplateGenerateContentRequest.swift b/FirebaseAI/Sources/TemplateGenerateContentRequest.swift new file mode 100644 index 00000000000..df0c442d495 --- /dev/null +++ b/FirebaseAI/Sources/TemplateGenerateContentRequest.swift @@ -0,0 +1,56 @@ +// Copyright 2025 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, watchOS 8.0, *) +struct TemplateGenerateContentRequest: Sendable { + let template: String + let variables: [String: TemplateVariable] + let history: [ModelContent] + let projectID: String + + let apiConfig: APIConfig + let options: RequestOptions +} + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +extension TemplateGenerateContentRequest: Encodable { + enum CodingKeys: String, CodingKey { + case variables = "inputs" + case history + } + + func encode(to encoder: any Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(variables, forKey: .variables) + try container.encode(history, forKey: .history) + } +} + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +extension TemplateGenerateContentRequest: GenerativeAIRequest { + typealias Response = GenerateContentResponse + + var url: URL { + var urlString = + "\(apiConfig.service.endpoint.rawValue)/\(apiConfig.version.rawValue)/projects/\(projectID)" + if case let .vertexAI(_, location) = apiConfig.service { + urlString += "/locations/\(location)" + } + let templateName = template.hasSuffix(".prompt") ? template : "\(template).prompt" + urlString += "/templates/\(templateName):templateGenerateContent" + return URL(string: urlString)! + } +} diff --git a/FirebaseAI/Sources/TemplateGenerativeModel.swift b/FirebaseAI/Sources/TemplateGenerativeModel.swift index 90df6d2d0f8..333ab8101dd 100644 --- a/FirebaseAI/Sources/TemplateGenerativeModel.swift +++ b/FirebaseAI/Sources/TemplateGenerativeModel.swift @@ -30,13 +30,15 @@ public final class TemplateGenerativeModel: BaseTemplateAPIClientModel { /// - variables: A dictionary of variables to substitute into the template. /// - Returns: The content generated by the model. /// - Throws: A ``GenerateContentError`` if the request failed. - public func templateGenerateContent(template: String, - variables: [String: Any]) async throws + public func generateContent(template: String, + variables: [String: Any], + options: RequestOptions = RequestOptions()) async throws -> GenerateContentResponse { - return try await templateGenerateContentWithHistory( + return try await generateContentWithHistory( history: [], template: template, - variables: variables + variables: variables, + options: options ) } @@ -48,16 +50,22 @@ public final class TemplateGenerativeModel: BaseTemplateAPIClientModel { /// - variables: A dictionary of variables to substitute into the template. /// - Returns: The content generated by the model. /// - Throws: A ``GenerateContentError`` if the request failed. - func templateGenerateContentWithHistory(history: [ModelContent], template: String, - variables: [String: Any]) async throws + func generateContentWithHistory(history: [ModelContent], template: String, + variables: [String: Any], + options: RequestOptions = RequestOptions()) async throws -> GenerateContentResponse { - let candidate = Candidate( - content: ModelContent(role: "model", parts: ["response"]), - safetyRatings: [], - finishReason: .stop, - citationMetadata: nil + let templateVariables = try variables.mapValues { try TemplateVariable(value: $0) } + let request = TemplateGenerateContentRequest( + template: template, + variables: templateVariables, + history: history, + projectID: generativeAIService.firebaseInfo.projectID, + apiConfig: apiConfig, + options: options ) - return GenerateContentResponse(candidates: [candidate], promptFeedback: nil) + let response: GenerateContentResponse = try await generativeAIService + .loadRequest(request: request) + return response } /// Creates a new chat conversation using this model with the provided history and template. @@ -66,10 +74,10 @@ public final class TemplateGenerativeModel: BaseTemplateAPIClientModel { /// - template: The prompt template to use. /// - history: The conversation history to use. /// - Returns: A new ``TemplateChatSession`` instance. - public func startTemplateChat(template: String, - history: [ModelContent] = []) -> TemplateChatSession { + public func startChat(template: String, + history: [ModelContent] = []) -> TemplateChatSession { return TemplateChatSession( - templateGenerateContent: templateGenerateContentWithHistory, + generateContent: generateContentWithHistory, template: template, history: history ) diff --git a/FirebaseAI/Sources/TemplateImagenModel.swift b/FirebaseAI/Sources/TemplateImagenModel.swift index ca93f9a74d2..f6f693fc09d 100644 --- a/FirebaseAI/Sources/TemplateImagenModel.swift +++ b/FirebaseAI/Sources/TemplateImagenModel.swift @@ -27,13 +27,14 @@ public final class TemplateImagenModel: BaseTemplateAPIClientModel { /// - variables: A dictionary of variables to substitute into the template. /// - Returns: The images generated by the model. /// - Throws: An error if the request failed. - public func templateImagenGenerateImages(template: String, - variables: [String: Any], - options: RequestOptions = RequestOptions()) async throws + public func generateImages(template: String, + variables: [String: Any], + options: RequestOptions = RequestOptions()) async throws -> GenerateImagesResponse { + let templateVariables = try variables.mapValues { try TemplateVariable(value: $0) } let request = GenerateImagesRequest( template: template, - variables: variables.mapValues { AnyCodable($0) }, + variables: templateVariables, apiConfig: apiConfig, options: options ) diff --git a/FirebaseAI/Sources/AnyCodable.swift b/FirebaseAI/Sources/TemplateVariable.swift similarity index 61% rename from FirebaseAI/Sources/AnyCodable.swift rename to FirebaseAI/Sources/TemplateVariable.swift index 43110196c4d..7ccdc381f83 100644 --- a/FirebaseAI/Sources/AnyCodable.swift +++ b/FirebaseAI/Sources/TemplateVariable.swift @@ -15,28 +15,28 @@ import Foundation -struct AnyCodable: Encodable { - let value: Any +enum TemplateVariable: Encodable, Sendable { + case string(String) + case int(Int) + case double(Double) + case bool(Bool) + case array([TemplateVariable]) + case dictionary([String: TemplateVariable]) - init(_ value: Any) { - self.value = value - } - - func encode(to encoder: Encoder) throws { - var container = encoder.singleValueContainer() + init(value: Any) throws { switch value { case let value as String: - try container.encode(value) + self = .string(value) case let value as Int: - try container.encode(value) + self = .int(value) case let value as Double: - try container.encode(value) + self = .double(value) case let value as Bool: - try container.encode(value) + self = .bool(value) case let value as [Any]: - try container.encode(value.map { AnyCodable($0) }) + self = try .array(value.map { try TemplateVariable(value: $0) }) case let value as [String: Any]: - try container.encode(value.mapValues { AnyCodable($0) }) + self = try .dictionary(value.mapValues { try TemplateVariable(value: $0) }) default: throw EncodingError.invalidValue( value, @@ -44,4 +44,22 @@ struct AnyCodable: Encodable { ) } } + + func encode(to encoder: Encoder) throws { + var container = encoder.singleValueContainer() + switch self { + case let .string(value): + try container.encode(value) + case let .int(value): + try container.encode(value) + case let .double(value): + try container.encode(value) + case let .bool(value): + try container.encode(value) + case let .array(value): + try container.encode(value) + case let .dictionary(value): + try container.encode(value) + } + } } diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift new file mode 100644 index 00000000000..446e87de026 --- /dev/null +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift @@ -0,0 +1,108 @@ + +// Copyright 2025 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 XCTest + +import FirebaseAI + +final class ServerPromptTemplateIntegrationTests: XCTestCase { + override func setUp() { + super.setUp() + continueAfterFailure = false + } + + func testGenerateContentWithText() async throws { + let model = FirebaseAI.firebaseAI(backend: .vertexAI()).templateGenerativeModel() + let userName = "paul" + do { + let response = try await model.generateContent( + template: "greeting", + variables: [ + "name": userName, + "language": "English", + ] + ) + let text = try XCTUnwrap(response.text) + print(text) + XCTAssert(text.contains("Paul")) + } catch { + XCTFail("An error occurred: \(error)") + } + } + + func testGenerateContentWithMedia() async throws { + let model = FirebaseAI.firebaseAI(backend: .googleAI()).templateGenerativeModel() + let image = UIImage(systemName: "photo")! + if let imageBytes = image.jpegData(compressionQuality: 0.8) { + let base64Image = imageBytes.base64EncodedString() + + do { + let response = try await model.generateContent( + template: "media", + variables: [ + "imageData": [ + "isInline": true, + "mimeType": "image/jpeg", + "contents": base64Image, + ], + ] + ) + XCTAssert(response.text?.isEmpty == false) + } catch { + XCTFail("An error occurred: \(error)") + } + } else { + XCTFail("Could not get image data.") + } + } + + func testGenerateImages() async throws { + let imagenModel = FirebaseAI.firebaseAI(backend: .googleAI()).templateImagenModel() + let imagenPrompt = "A cat picture" + do { + let response = try await imagenModel.generateImages( + template: "generate_images", + variables: [ + "prompt": imagenPrompt, + ] + ) + XCTAssertEqual(response.images.count, 1) + } catch { + XCTFail("An error occurred: \(error)") + } + } + + func testChat() async throws { + let model = FirebaseAI.firebaseAI(backend: .googleAI()).templateGenerativeModel() + let initialHistory = [ + ModelContent(role: "user", parts: "Hello!"), + ModelContent(role: "model", parts: "Hi there! How can I help?"), + ] + let chatSession = model.startChat(template: "chat_history", history: initialHistory) + + let userMessage = "What's the weather like?" + + do { + let response = try await chatSession.sendMessage( + userMessage, + variables: ["message": userMessage] + ) + XCTAssert(response.text?.isEmpty == false) + XCTAssertEqual(chatSession.history.count, 3) + } catch { + XCTFail("An error occurred: \(error)") + } + } +} diff --git a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift index d3cdff52168..15737cce7df 100644 --- a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift @@ -32,7 +32,7 @@ final class TemplateChatSessionTests: XCTestCase { } func testSendMessage() async throws { - let chat = model.startTemplateChat(template: "test-template") + let chat = model.startChat(template: "test-template") let response = try await chat.sendMessage("Hello", variables: ["name": "test"]) XCTAssertEqual(chat.history.count, 2) XCTAssertEqual(chat.history[0].role, "user") diff --git a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift index 2abd1f3d685..5dcc857722d 100644 --- a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift @@ -30,11 +30,11 @@ final class TemplateImagenModelTests: XCTestCase { firebaseInfo: firebaseInfo, urlSession: urlSession ) - let apiConfig = APIConfig(service: .googleAI(endpoint: .firebaseProxyProd), version: .v1beta) + let apiConfig = APIConfig(service: .googleAI(endpoint: .firebaseProxyProd), version: .v1beta) model = TemplateImagenModel(generativeAIService: generativeAIService, apiConfig: apiConfig) } - func testTemplateImagenGenerateImages() async throws { + func testGenerateImages() async throws { MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( forResource: "unary-success-image-response", withExtension: "json", @@ -42,7 +42,7 @@ final class TemplateImagenModelTests: XCTestCase { isImageRequest: true ) - let response = try await model.templateImagenGenerateImages( + let response = try await model.generateImages( template: "test-template", variables: ["prompt": "a cat picture"] ) From 8edb1013644c46d05c74fef07219b9e4e48a377a Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Thu, 9 Oct 2025 15:45:53 -0700 Subject: [PATCH 05/24] FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift --- .../TemplateGenerateContentRequest.swift | 1 - .../Unit/TemplateGenerativeModelTests.swift | 52 +++++++++++++++++++ .../Tests/Unit/TemplateImagenModelTests.swift | 4 +- .../GenerativeModelTestUtil.swift | 6 ++- 4 files changed, 58 insertions(+), 5 deletions(-) create mode 100644 FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift diff --git a/FirebaseAI/Sources/TemplateGenerateContentRequest.swift b/FirebaseAI/Sources/TemplateGenerateContentRequest.swift index df0c442d495..578571672a9 100644 --- a/FirebaseAI/Sources/TemplateGenerateContentRequest.swift +++ b/FirebaseAI/Sources/TemplateGenerateContentRequest.swift @@ -20,7 +20,6 @@ struct TemplateGenerateContentRequest: Sendable { let variables: [String: TemplateVariable] let history: [ModelContent] let projectID: String - let apiConfig: APIConfig let options: RequestOptions } diff --git a/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift b/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift new file mode 100644 index 00000000000..0067bde27d0 --- /dev/null +++ b/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift @@ -0,0 +1,52 @@ + +// Copyright 2025 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. + +@testable import FirebaseAI +import FirebaseCore +import XCTest + +final class TemplateGenerativeModelTests: XCTestCase { + var urlSession: URLSession! + var model: TemplateGenerativeModel! + + override func setUp() { + super.setUp() + let configuration = URLSessionConfiguration.default + configuration.protocolClasses = [MockURLProtocol.self] + urlSession = URLSession(configuration: configuration) + let firebaseInfo = GenerativeModelTestUtil.testFirebaseInfo() + let generativeAIService = GenerativeAIService( + firebaseInfo: firebaseInfo, + urlSession: urlSession + ) + let apiConfig = APIConfig(service: .googleAI(endpoint: .firebaseProxyProd), version: .v1beta) + model = TemplateGenerativeModel(generativeAIService: generativeAIService, apiConfig: apiConfig) + } + + func testGenerateContent() async throws { + MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( + forResource: "unary-success-response", + withExtension: "json", + subdirectory: "mock-responses", + isTemplateRequest: true + ) + + let response = try await model.generateContent( + template: "test-template", + variables: ["name": "test"] + ) + XCTAssertEqual(response.text, "Hello, world!") + } +} diff --git a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift index 5dcc857722d..ec8e0809721 100644 --- a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift @@ -38,8 +38,8 @@ final class TemplateImagenModelTests: XCTestCase { MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( forResource: "unary-success-image-response", withExtension: "json", - subdirectory: "vertexai-sdk-test-data/mock-responses", - isImageRequest: true + subdirectory: "mock-responses", + isTemplateRequest: true ) let response = try await model.generateImages( diff --git a/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift b/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift index aa2a67c2d34..39c43ae0fc2 100644 --- a/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift +++ b/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift @@ -31,7 +31,7 @@ enum GenerativeModelTestUtil { appCheckToken: String? = nil, authToken: String? = nil, dataCollection: Bool = true, - isImageRequest: Bool = false) throws -> ((URLRequest) throws -> ( + isTemplateRequest: Bool = false) throws -> ((URLRequest) throws -> ( URLResponse, AsyncLineSequence? )) { @@ -46,7 +46,9 @@ enum GenerativeModelTestUtil { ) return { request in let requestURL = try XCTUnwrap(request.url) - if !isImageRequest { + if isTemplateRequest { + XCTAssertEqual(requestURL.path.occurrenceCount(of: "templates/test-template.prompt:template"), 1) + } else { XCTAssertEqual(requestURL.path.occurrenceCount(of: "models/"), 1) } XCTAssertEqual(request.timeoutInterval, timeout) From 1a47800cf04e6d7cdf725c8f7a817e007d0b1387 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Fri, 10 Oct 2025 11:07:38 -0700 Subject: [PATCH 06/24] CI fixes --- .../Tests/Unit/TemplateChatSessionTests.swift | 1 + .../Tests/Unit/TemplateGenerativeModelTests.swift | 1 + .../Tests/Unit/TemplateImagenModelTests.swift | 1 + .../TestUtilities/GenerativeModelTestUtil.swift | 14 +++++++++----- 4 files changed, 12 insertions(+), 5 deletions(-) diff --git a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift index 15737cce7df..4bd920c0caf 100644 --- a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift @@ -17,6 +17,7 @@ import FirebaseCore import XCTest +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) final class TemplateChatSessionTests: XCTestCase { var model: TemplateGenerativeModel! diff --git a/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift b/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift index 0067bde27d0..c5928fb7aa6 100644 --- a/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift @@ -17,6 +17,7 @@ import FirebaseCore import XCTest +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) final class TemplateGenerativeModelTests: XCTestCase { var urlSession: URLSession! var model: TemplateGenerativeModel! diff --git a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift index ec8e0809721..1456bd4cea8 100644 --- a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift @@ -16,6 +16,7 @@ @testable import FirebaseAI import XCTest +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) final class TemplateImagenModelTests: XCTestCase { var urlSession: URLSession! var model: TemplateImagenModel! diff --git a/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift b/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift index 39c43ae0fc2..e783fdd6388 100644 --- a/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift +++ b/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift @@ -31,10 +31,11 @@ enum GenerativeModelTestUtil { appCheckToken: String? = nil, authToken: String? = nil, dataCollection: Bool = true, - isTemplateRequest: Bool = false) throws -> ((URLRequest) throws -> ( - URLResponse, - AsyncLineSequence? - )) { + isTemplateRequest: Bool = false) throws + -> ((URLRequest) throws -> ( + URLResponse, + AsyncLineSequence? + )) { // Skip tests using MockURLProtocol on watchOS; unsupported in watchOS 2 and later, see // https://developer.apple.com/documentation/foundation/urlprotocol for details. #if os(watchOS) @@ -47,7 +48,10 @@ enum GenerativeModelTestUtil { return { request in let requestURL = try XCTUnwrap(request.url) if isTemplateRequest { - XCTAssertEqual(requestURL.path.occurrenceCount(of: "templates/test-template.prompt:template"), 1) + XCTAssertEqual( + requestURL.path.occurrenceCount(of: "templates/test-template.prompt:template"), + 1 + ) } else { XCTAssertEqual(requestURL.path.occurrenceCount(of: "models/"), 1) } From 92db0ef810fcb8098b5f6873261d1f06fc8628ea Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Fri, 10 Oct 2025 12:13:37 -0700 Subject: [PATCH 07/24] Streaming APIs --- .../Sources/BaseTemplateAPIClientModel.swift | 28 ---- FirebaseAI/Sources/TemplateChatSession.swift | 139 ++++++++++++++++-- .../TemplateGenerateContentRequest.swift | 6 + .../Sources/TemplateGenerativeModel.swift | 54 ++++++- FirebaseAI/Sources/TemplateImagenModel.swift | 10 +- 5 files changed, 187 insertions(+), 50 deletions(-) delete mode 100644 FirebaseAI/Sources/BaseTemplateAPIClientModel.swift diff --git a/FirebaseAI/Sources/BaseTemplateAPIClientModel.swift b/FirebaseAI/Sources/BaseTemplateAPIClientModel.swift deleted file mode 100644 index 6363310ce5b..00000000000 --- a/FirebaseAI/Sources/BaseTemplateAPIClientModel.swift +++ /dev/null @@ -1,28 +0,0 @@ - -// Copyright 2025 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 - -/// A base class for template API client models. -@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) -public class BaseTemplateAPIClientModel { - let generativeAIService: GenerativeAIService - let apiConfig: APIConfig - - init(generativeAIService: GenerativeAIService, apiConfig: APIConfig) { - self.generativeAIService = generativeAIService - self.apiConfig = apiConfig - } -} diff --git a/FirebaseAI/Sources/TemplateChatSession.swift b/FirebaseAI/Sources/TemplateChatSession.swift index cbb4da9cb2d..7d2bf815774 100644 --- a/FirebaseAI/Sources/TemplateChatSession.swift +++ b/FirebaseAI/Sources/TemplateChatSession.swift @@ -1,4 +1,3 @@ - // Copyright 2024 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); @@ -17,31 +16,145 @@ import Foundation /// A chat session that allows for conversation with a model. @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) -public class TemplateChatSession { - private let generateContent: ([ModelContent], String, [String: Any], RequestOptions) async throws - -> GenerateContentResponse +public final class TemplateChatSession: Sendable { + private let model: TemplateGenerativeModel private let template: String - public var history: [ModelContent] - init(generateContent: @escaping (([ModelContent], String, [String: Any], - RequestOptions) async throws - -> GenerateContentResponse), - template: String, history: [ModelContent]) { - self.generateContent = generateContent + private let historyLock = NSLock() + private nonisolated(unsafe) var _history: [ModelContent] = [] + public var history: [ModelContent] { + get { + historyLock.withLock { _history } + } + set { + historyLock.withLock { _history = newValue } + } + } + + init(model: TemplateGenerativeModel, template: String, history: [ModelContent]) { + self.model = model self.template = template self.history = history } + private func appendHistory(contentsOf: [ModelContent]) { + historyLock.withLock { + _history.append(contentsOf: contentsOf) + } + } + + private func appendHistory(_ newElement: ModelContent) { + historyLock.withLock { + _history.append(newElement) + } + } + /// Sends a message to the model and returns the response. public func sendMessage(_ message: any PartsRepresentable, variables: [String: Any], options: RequestOptions = RequestOptions()) async throws -> GenerateContentResponse { - let response = try await generateContent(history, template, variables, options) - history.append(ModelContent(role: "user", parts: message.partsValue)) + let templateVariables = try variables.mapValues { try TemplateVariable(value: $0) } + let response = try await model.generateContentWithHistory( + history: history, + template: template, + variables: templateVariables, + options: options + ) + appendHistory(ModelContent(role: "user", parts: message.partsValue)) if let modelResponse = response.candidates.first { - history.append(modelResponse.content) + appendHistory(modelResponse.content) } return response } + + public func sendMessageStream(_ message: any PartsRepresentable, + variables: [String: Any], + options: RequestOptions = RequestOptions()) throws + -> AsyncThrowingStream { + let templateVariables = try variables.mapValues { try TemplateVariable(value: $0) } + let newContent = [ModelContent(role: "user", parts: message.partsValue)] + let stream = try model.generateContentStreamWithHistory( + history: history, + template: template, + variables: templateVariables, + options: options + ) + return AsyncThrowingStream { continuation in + Task { + var aggregatedContent: [ModelContent] = [] + + do { + for try await chunk in stream { + // Capture any content that's streaming. This should be populated if there's no error. + if let chunkContent = chunk.candidates.first?.content { + aggregatedContent.append(chunkContent) + } + + // Pass along the chunk. + continuation.yield(chunk) + } + } catch { + // Rethrow the error that the underlying stream threw. Don't add anything to history. + continuation.finish(throwing: error) + return + } + + // Save the request. + appendHistory(contentsOf: newContent) + + // Aggregate the content to add it to the history before we finish. + let aggregated = aggregatedChunks(aggregatedContent) + appendHistory(aggregated) + continuation.finish() + } + } + } + + private func aggregatedChunks(_ chunks: [ModelContent]) -> ModelContent { + var parts: [InternalPart] = [] + var combinedText = "" + var combinedThoughts = "" + + func flush() { + if !combinedThoughts.isEmpty { + parts.append(InternalPart(.text(combinedThoughts), isThought: true, thoughtSignature: nil)) + combinedThoughts = "" + } + if !combinedText.isEmpty { + parts.append(InternalPart(.text(combinedText), isThought: nil, thoughtSignature: nil)) + combinedText = "" + } + } + + // Loop through all the parts, aggregating the text. + for part in chunks.flatMap({ $0.internalParts }) { + // Only text parts may be combined. + if case let .text(text) = part.data, part.thoughtSignature == nil { + // Thought summaries must not be combined with regular text. + if part.isThought ?? false { + // If we were combining regular text, flush it before handling "thoughts". + if !combinedText.isEmpty { + flush() + } + combinedThoughts += text + } else { + // If we were combining "thoughts", flush it before handling regular text. + if !combinedThoughts.isEmpty { + flush() + } + combinedText += text + } + } else { + // This is a non-combinable part (not text), flush any pending text. + flush() + parts.append(part) + } + } + + // Flush any remaining text. + flush() + + return ModelContent(role: "model", parts: parts) + } } diff --git a/FirebaseAI/Sources/TemplateGenerateContentRequest.swift b/FirebaseAI/Sources/TemplateGenerateContentRequest.swift index 578571672a9..7c2da089b80 100644 --- a/FirebaseAI/Sources/TemplateGenerateContentRequest.swift +++ b/FirebaseAI/Sources/TemplateGenerateContentRequest.swift @@ -20,6 +20,8 @@ struct TemplateGenerateContentRequest: Sendable { let variables: [String: TemplateVariable] let history: [ModelContent] let projectID: String + let stream: Bool + let apiConfig: APIConfig let options: RequestOptions } @@ -50,6 +52,10 @@ extension TemplateGenerateContentRequest: GenerativeAIRequest { } let templateName = template.hasSuffix(".prompt") ? template : "\(template).prompt" urlString += "/templates/\(templateName):templateGenerateContent" + if stream { + // TODO: Fix this. + urlString += "?alt=sse" + } return URL(string: urlString)! } } diff --git a/FirebaseAI/Sources/TemplateGenerativeModel.swift b/FirebaseAI/Sources/TemplateGenerativeModel.swift index 333ab8101dd..98df95fba9c 100644 --- a/FirebaseAI/Sources/TemplateGenerativeModel.swift +++ b/FirebaseAI/Sources/TemplateGenerativeModel.swift @@ -18,9 +18,13 @@ import Foundation /// A type that represents a remote multimodal model (like Gemini), with the ability to generate /// content based on various input types. @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) -public final class TemplateGenerativeModel: BaseTemplateAPIClientModel { - override init(generativeAIService: GenerativeAIService, apiConfig: APIConfig) { - super.init(generativeAIService: generativeAIService, apiConfig: apiConfig) +public final class TemplateGenerativeModel: Sendable { + let generativeAIService: GenerativeAIService + let apiConfig: APIConfig + + init(generativeAIService: GenerativeAIService, apiConfig: APIConfig) { + self.generativeAIService = generativeAIService + self.apiConfig = apiConfig } /// Generates content from a prompt template and variables. @@ -34,10 +38,11 @@ public final class TemplateGenerativeModel: BaseTemplateAPIClientModel { variables: [String: Any], options: RequestOptions = RequestOptions()) async throws -> GenerateContentResponse { + let templateVariables = try variables.mapValues { try TemplateVariable(value: $0) } return try await generateContentWithHistory( history: [], template: template, - variables: variables, + variables: templateVariables, options: options ) } @@ -51,15 +56,15 @@ public final class TemplateGenerativeModel: BaseTemplateAPIClientModel { /// - Returns: The content generated by the model. /// - Throws: A ``GenerateContentError`` if the request failed. func generateContentWithHistory(history: [ModelContent], template: String, - variables: [String: Any], + variables: [String: TemplateVariable], options: RequestOptions = RequestOptions()) async throws -> GenerateContentResponse { - let templateVariables = try variables.mapValues { try TemplateVariable(value: $0) } let request = TemplateGenerateContentRequest( template: template, - variables: templateVariables, + variables: variables, history: history, projectID: generativeAIService.firebaseInfo.projectID, + stream: false, apiConfig: apiConfig, options: options ) @@ -68,6 +73,39 @@ public final class TemplateGenerativeModel: BaseTemplateAPIClientModel { return response } + public func generateContentStream(template: String, + variables: [String: Any], + options: RequestOptions = RequestOptions()) throws + -> AsyncThrowingStream { + let templateVariables = try variables.mapValues { try TemplateVariable(value: $0) } + let request = TemplateGenerateContentRequest( + template: template, + variables: templateVariables, + history: [], + projectID: generativeAIService.firebaseInfo.projectID, + stream: true, + apiConfig: apiConfig, + options: options + ) + return generativeAIService.loadRequestStream(request: request) + } + + func generateContentStreamWithHistory(history: [ModelContent], template: String, + variables: [String: TemplateVariable], + options: RequestOptions = RequestOptions()) throws + -> AsyncThrowingStream { + let request = TemplateGenerateContentRequest( + template: template, + variables: variables, + history: history, + projectID: generativeAIService.firebaseInfo.projectID, + stream: true, + apiConfig: apiConfig, + options: options + ) + return generativeAIService.loadRequestStream(request: request) + } + /// Creates a new chat conversation using this model with the provided history and template. /// /// - Parameters: @@ -77,7 +115,7 @@ public final class TemplateGenerativeModel: BaseTemplateAPIClientModel { public func startChat(template: String, history: [ModelContent] = []) -> TemplateChatSession { return TemplateChatSession( - generateContent: generateContentWithHistory, + model: self, template: template, history: history ) diff --git a/FirebaseAI/Sources/TemplateImagenModel.swift b/FirebaseAI/Sources/TemplateImagenModel.swift index f6f693fc09d..406dde67aa7 100644 --- a/FirebaseAI/Sources/TemplateImagenModel.swift +++ b/FirebaseAI/Sources/TemplateImagenModel.swift @@ -19,7 +19,15 @@ import Foundation /// generate /// images based on various input types. @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) -public final class TemplateImagenModel: BaseTemplateAPIClientModel { +public final class TemplateImagenModel: Sendable { + let generativeAIService: GenerativeAIService + let apiConfig: APIConfig + + init(generativeAIService: GenerativeAIService, apiConfig: APIConfig) { + self.generativeAIService = generativeAIService + self.apiConfig = apiConfig + } + /// Generates images from a prompt template and variables. /// /// - Parameters: From 004fb94aca47f96febf66a264d540d138b270057 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Fri, 10 Oct 2025 12:36:49 -0700 Subject: [PATCH 08/24] chat history refactor --- FirebaseAI/Sources/Chat.swift | 79 ++-------------- FirebaseAI/Sources/History.swift | 95 ++++++++++++++++++++ FirebaseAI/Sources/TemplateChatSession.swift | 93 +++++-------------- 3 files changed, 125 insertions(+), 142 deletions(-) create mode 100644 FirebaseAI/Sources/History.swift diff --git a/FirebaseAI/Sources/Chat.swift b/FirebaseAI/Sources/Chat.swift index 80e908a8f57..99c6fb13367 100644 --- a/FirebaseAI/Sources/Chat.swift +++ b/FirebaseAI/Sources/Chat.swift @@ -19,35 +19,21 @@ import Foundation @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) public final class Chat: Sendable { private let model: GenerativeModel + private let _history: History - /// Initializes a new chat representing a 1:1 conversation between model and user. init(model: GenerativeModel, history: [ModelContent]) { self.model = model - self.history = history + _history = History(history: history) } - private let historyLock = NSLock() - private nonisolated(unsafe) var _history: [ModelContent] = [] /// The previous content from the chat that has been successfully sent and received from the /// model. This will be provided to the model for each message sent as context for the discussion. public var history: [ModelContent] { get { - historyLock.withLock { _history } + return _history.history } set { - historyLock.withLock { _history = newValue } - } - } - - private func appendHistory(contentsOf: [ModelContent]) { - historyLock.withLock { - _history.append(contentsOf: contentsOf) - } - } - - private func appendHistory(_ newElement: ModelContent) { - historyLock.withLock { - _history.append(newElement) + _history.history = newValue } } @@ -87,8 +73,8 @@ public final class Chat: Sendable { let toAdd = ModelContent(role: "model", parts: reply.parts) // Append the request and successful result to history, then return the value. - appendHistory(contentsOf: newContent) - appendHistory(toAdd) + _history.append(contentsOf: newContent) + _history.append(toAdd) return result } @@ -136,63 +122,16 @@ public final class Chat: Sendable { } // Save the request. - appendHistory(contentsOf: newContent) + _history.append(contentsOf: newContent) // Aggregate the content to add it to the history before we finish. - let aggregated = self.aggregatedChunks(aggregatedContent) - self.appendHistory(aggregated) + let aggregated = self._history.aggregatedChunks(aggregatedContent) + self._history.append(aggregated) continuation.finish() } } } - private func aggregatedChunks(_ chunks: [ModelContent]) -> ModelContent { - var parts: [InternalPart] = [] - var combinedText = "" - var combinedThoughts = "" - - func flush() { - if !combinedThoughts.isEmpty { - parts.append(InternalPart(.text(combinedThoughts), isThought: true, thoughtSignature: nil)) - combinedThoughts = "" - } - if !combinedText.isEmpty { - parts.append(InternalPart(.text(combinedText), isThought: nil, thoughtSignature: nil)) - combinedText = "" - } - } - - // Loop through all the parts, aggregating the text. - for part in chunks.flatMap({ $0.internalParts }) { - // Only text parts may be combined. - if case let .text(text) = part.data, part.thoughtSignature == nil { - // Thought summaries must not be combined with regular text. - if part.isThought ?? false { - // If we were combining regular text, flush it before handling "thoughts". - if !combinedText.isEmpty { - flush() - } - combinedThoughts += text - } else { - // If we were combining "thoughts", flush it before handling regular text. - if !combinedThoughts.isEmpty { - flush() - } - combinedText += text - } - } else { - // This is a non-combinable part (not text), flush any pending text. - flush() - parts.append(part) - } - } - - // Flush any remaining text. - flush() - - return ModelContent(role: "model", parts: parts) - } - /// Populates the `role` field with `user` if it doesn't exist. Required in chat sessions. private func populateContentRole(_ content: ModelContent) -> ModelContent { if content.role != nil { diff --git a/FirebaseAI/Sources/History.swift b/FirebaseAI/Sources/History.swift new file mode 100644 index 00000000000..7e5b81ad60b --- /dev/null +++ b/FirebaseAI/Sources/History.swift @@ -0,0 +1,95 @@ + +// Copyright 2024 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, watchOS 8.0, *) +final class History: Sendable { + private let historyLock = NSLock() + private nonisolated(unsafe) var _history: [ModelContent] = [] + /// The previous content from the chat that has been successfully sent and received from the + /// model. This will be provided to the model for each message sent as context for the discussion. + public var history: [ModelContent] { + get { + historyLock.withLock { _history } + } + set { + historyLock.withLock { _history = newValue } + } + } + + init(history: [ModelContent]) { + self.history = history + } + + func append(contentsOf: [ModelContent]) { + historyLock.withLock { + _history.append(contentsOf: contentsOf) + } + } + + func append(_ newElement: ModelContent) { + historyLock.withLock { + _history.append(newElement) + } + } + + func aggregatedChunks(_ chunks: [ModelContent]) -> ModelContent { + var parts: [InternalPart] = [] + var combinedText = "" + var combinedThoughts = "" + + func flush() { + if !combinedThoughts.isEmpty { + parts.append(InternalPart(.text(combinedThoughts), isThought: true, thoughtSignature: nil)) + combinedThoughts = "" + } + if !combinedText.isEmpty { + parts.append(InternalPart(.text(combinedText), isThought: nil, thoughtSignature: nil)) + combinedText = "" + } + } + + // Loop through all the parts, aggregating the text. + for part in chunks.flatMap({ $0.internalParts }) { + // Only text parts may be combined. + if case let .text(text) = part.data, part.thoughtSignature == nil { + // Thought summaries must not be combined with regular text. + if part.isThought ?? false { + // If we were combining regular text, flush it before handling "thoughts". + if !combinedText.isEmpty { + flush() + } + combinedThoughts += text + } else { + // If we were combining "thoughts", flush it before handling regular text. + if !combinedThoughts.isEmpty { + flush() + } + combinedText += text + } + } else { + // This is a non-combinable part (not text), flush any pending text. + flush() + parts.append(part) + } + } + + // Flush any remaining text. + flush() + + return ModelContent(role: "model", parts: parts) + } +} diff --git a/FirebaseAI/Sources/TemplateChatSession.swift b/FirebaseAI/Sources/TemplateChatSession.swift index 7d2bf815774..fabd4fcd52b 100644 --- a/FirebaseAI/Sources/TemplateChatSession.swift +++ b/FirebaseAI/Sources/TemplateChatSession.swift @@ -19,33 +19,20 @@ import Foundation public final class TemplateChatSession: Sendable { private let model: TemplateGenerativeModel private let template: String - - private let historyLock = NSLock() - private nonisolated(unsafe) var _history: [ModelContent] = [] - public var history: [ModelContent] { - get { - historyLock.withLock { _history } - } - set { - historyLock.withLock { _history = newValue } - } - } + private let _history: History init(model: TemplateGenerativeModel, template: String, history: [ModelContent]) { self.model = model self.template = template - self.history = history + _history = History(history: history) } - private func appendHistory(contentsOf: [ModelContent]) { - historyLock.withLock { - _history.append(contentsOf: contentsOf) + public var history: [ModelContent] { + get { + return _history.history } - } - - private func appendHistory(_ newElement: ModelContent) { - historyLock.withLock { - _history.append(newElement) + set { + _history.history = newValue } } @@ -55,15 +42,16 @@ public final class TemplateChatSession: Sendable { options: RequestOptions = RequestOptions()) async throws -> GenerateContentResponse { let templateVariables = try variables.mapValues { try TemplateVariable(value: $0) } + let newContent = populateContentRole(ModelContent(parts: message.partsValue)) let response = try await model.generateContentWithHistory( - history: history, + history: _history.history + [newContent], template: template, variables: templateVariables, options: options ) - appendHistory(ModelContent(role: "user", parts: message.partsValue)) + _history.append(newContent) if let modelResponse = response.candidates.first { - appendHistory(modelResponse.content) + _history.append(modelResponse.content) } return response } @@ -73,9 +61,9 @@ public final class TemplateChatSession: Sendable { options: RequestOptions = RequestOptions()) throws -> AsyncThrowingStream { let templateVariables = try variables.mapValues { try TemplateVariable(value: $0) } - let newContent = [ModelContent(role: "user", parts: message.partsValue)] + let newContent = populateContentRole(ModelContent(parts: message.partsValue)) let stream = try model.generateContentStreamWithHistory( - history: history, + history: _history.history + [newContent], template: template, variables: templateVariables, options: options @@ -101,60 +89,21 @@ public final class TemplateChatSession: Sendable { } // Save the request. - appendHistory(contentsOf: newContent) + _history.append(newContent) // Aggregate the content to add it to the history before we finish. - let aggregated = aggregatedChunks(aggregatedContent) - appendHistory(aggregated) + let aggregated = _history.aggregatedChunks(aggregatedContent) + _history.append(aggregated) continuation.finish() } } } - private func aggregatedChunks(_ chunks: [ModelContent]) -> ModelContent { - var parts: [InternalPart] = [] - var combinedText = "" - var combinedThoughts = "" - - func flush() { - if !combinedThoughts.isEmpty { - parts.append(InternalPart(.text(combinedThoughts), isThought: true, thoughtSignature: nil)) - combinedThoughts = "" - } - if !combinedText.isEmpty { - parts.append(InternalPart(.text(combinedText), isThought: nil, thoughtSignature: nil)) - combinedText = "" - } + private func populateContentRole(_ content: ModelContent) -> ModelContent { + if content.role != nil { + return content + } else { + return ModelContent(role: "user", parts: content.parts) } - - // Loop through all the parts, aggregating the text. - for part in chunks.flatMap({ $0.internalParts }) { - // Only text parts may be combined. - if case let .text(text) = part.data, part.thoughtSignature == nil { - // Thought summaries must not be combined with regular text. - if part.isThought ?? false { - // If we were combining regular text, flush it before handling "thoughts". - if !combinedText.isEmpty { - flush() - } - combinedThoughts += text - } else { - // If we were combining "thoughts", flush it before handling regular text. - if !combinedThoughts.isEmpty { - flush() - } - combinedText += text - } - } else { - // This is a non-combinable part (not text), flush any pending text. - flush() - parts.append(part) - } - } - - // Flush any remaining text. - flush() - - return ModelContent(role: "model", parts: parts) } } From a475dba9dccc83279b846f8d8768b830ae6e3626 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Fri, 10 Oct 2025 14:48:42 -0700 Subject: [PATCH 09/24] Re-sort tests --- ...ServerPromptTemplateIntegrationTests.swift | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift index 446e87de026..86fdcfcb3cf 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift @@ -42,6 +42,22 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { } } + func testGenerateImages() async throws { + let imagenModel = FirebaseAI.firebaseAI(backend: .googleAI()).templateImagenModel() + let imagenPrompt = "A cat picture" + do { + let response = try await imagenModel.generateImages( + template: "describeImages", + variables: [ + "prompt": imagenPrompt, + ] + ) + XCTAssertEqual(response.images.count, 1) + } catch { + XCTFail("An error occurred: \(error)") + } + } + func testGenerateContentWithMedia() async throws { let model = FirebaseAI.firebaseAI(backend: .googleAI()).templateGenerativeModel() let image = UIImage(systemName: "photo")! @@ -68,22 +84,6 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { } } - func testGenerateImages() async throws { - let imagenModel = FirebaseAI.firebaseAI(backend: .googleAI()).templateImagenModel() - let imagenPrompt = "A cat picture" - do { - let response = try await imagenModel.generateImages( - template: "generate_images", - variables: [ - "prompt": imagenPrompt, - ] - ) - XCTAssertEqual(response.images.count, 1) - } catch { - XCTFail("An error occurred: \(error)") - } - } - func testChat() async throws { let model = FirebaseAI.firebaseAI(backend: .googleAI()).templateGenerativeModel() let initialHistory = [ From 2f718d896f2c4f525ad3a6f02620179f1e6b1f88 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Fri, 10 Oct 2025 21:55:30 +0000 Subject: [PATCH 10/24] Fix TemplateImagenModel to encode variables as inputs Updated `GenerateImagesRequest` to encode the `variables` property as `inputs`, similar to how `TemplateGenerateContentRequest` handles it. This ensures consistency across the API. --- FirebaseAI/Sources/GenerateImagesRequest.swift | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/FirebaseAI/Sources/GenerateImagesRequest.swift b/FirebaseAI/Sources/GenerateImagesRequest.swift index 81fe0f5dde8..b3953c2e6ef 100644 --- a/FirebaseAI/Sources/GenerateImagesRequest.swift +++ b/FirebaseAI/Sources/GenerateImagesRequest.swift @@ -39,13 +39,11 @@ public class GenerateImagesRequest: @unchecked Sendable, GenerativeAIRequest { } enum CodingKeys: String, CodingKey { - case template - case variables + case variables = "inputs" } public func encode(to encoder: Encoder) throws { var container = encoder.container(keyedBy: CodingKeys.self) - try container.encode(template, forKey: .template) try container.encode(variables, forKey: .variables) } } From cca457a77bcda57242a6dc785169323b7dcf436f Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Fri, 10 Oct 2025 22:03:32 +0000 Subject: [PATCH 11/24] Fix: Correct URL construction for image generation requests Updated `GenerateImagesRequest` to correctly construct the URL, including the `projectID` and `location`. This mirrors the URL construction in `TemplateGenerateContentRequest` and fixes a 404 error that was occurring in the `testGenerateImages` integration test. The `TemplateImagenModel` was also updated to pass the `projectID` to the `GenerateImagesRequest` initializer. --- .../Sources/GenerateImagesRequest.swift | 20 +++++++++++++------ FirebaseAI/Sources/TemplateImagenModel.swift | 2 ++ 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/FirebaseAI/Sources/GenerateImagesRequest.swift b/FirebaseAI/Sources/GenerateImagesRequest.swift index b3953c2e6ef..f0dc56e7935 100644 --- a/FirebaseAI/Sources/GenerateImagesRequest.swift +++ b/FirebaseAI/Sources/GenerateImagesRequest.swift @@ -19,23 +19,31 @@ import Foundation public class GenerateImagesRequest: @unchecked Sendable, GenerativeAIRequest { public typealias Response = GenerateImagesResponse - public let url: URL + public var url: URL { + var urlString = + "\(apiConfig.service.endpoint.rawValue)/\(apiConfig.version.rawValue)/projects/\(projectID)" + if case let .vertexAI(_, location) = apiConfig.service { + urlString += "/locations/\(location)" + } + urlString += "/templates/\(template):\(ImageAPIMethod.generateImages.rawValue)" + return URL(string: urlString)! + } + public let options: RequestOptions let apiConfig: APIConfig let template: String let variables: [String: TemplateVariable] + let projectID: String - init(template: String, variables: [String: TemplateVariable], apiConfig: APIConfig, - options: RequestOptions) { - let modelURL = - "\(apiConfig.service.endpoint.rawValue)/\(apiConfig.version.rawValue)/\(template)" - url = URL(string: "\(modelURL):\(ImageAPIMethod.generateImages.rawValue)")! + init(template: String, variables: [String: TemplateVariable], projectID: String, + apiConfig: APIConfig, options: RequestOptions) { self.apiConfig = apiConfig self.options = options self.template = template self.variables = variables + self.projectID = projectID } enum CodingKeys: String, CodingKey { diff --git a/FirebaseAI/Sources/TemplateImagenModel.swift b/FirebaseAI/Sources/TemplateImagenModel.swift index 406dde67aa7..3c8efc31425 100644 --- a/FirebaseAI/Sources/TemplateImagenModel.swift +++ b/FirebaseAI/Sources/TemplateImagenModel.swift @@ -40,9 +40,11 @@ public final class TemplateImagenModel: Sendable { options: RequestOptions = RequestOptions()) async throws -> GenerateImagesResponse { let templateVariables = try variables.mapValues { try TemplateVariable(value: $0) } + let projectID = generativeAIService.firebaseInfo.projectID let request = GenerateImagesRequest( template: template, variables: templateVariables, + projectID: projectID, apiConfig: apiConfig, options: options ) From 2b2b15f8084384fe48a231882012dae22138f005 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Sat, 11 Oct 2025 20:04:15 -0700 Subject: [PATCH 12/24] checkpoint --- FirebaseAI/Sources/GenerateImagesRequest.swift | 3 ++- FirebaseAI/Sources/ImageAPIMethod.swift | 2 +- .../Integration/ServerPromptTemplateIntegrationTests.swift | 4 ++-- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/FirebaseAI/Sources/GenerateImagesRequest.swift b/FirebaseAI/Sources/GenerateImagesRequest.swift index f0dc56e7935..4343589bdb2 100644 --- a/FirebaseAI/Sources/GenerateImagesRequest.swift +++ b/FirebaseAI/Sources/GenerateImagesRequest.swift @@ -25,7 +25,8 @@ public class GenerateImagesRequest: @unchecked Sendable, GenerativeAIRequest { if case let .vertexAI(_, location) = apiConfig.service { urlString += "/locations/\(location)" } - urlString += "/templates/\(template):\(ImageAPIMethod.generateImages.rawValue)" + let templateName = template.hasSuffix(".prompt") ? template : "\(template).prompt" + urlString += "/templates/\(templateName):\(ImageAPIMethod.generateImages.rawValue)" return URL(string: urlString)! } diff --git a/FirebaseAI/Sources/ImageAPIMethod.swift b/FirebaseAI/Sources/ImageAPIMethod.swift index 84afa50a742..7a5efbaa032 100644 --- a/FirebaseAI/Sources/ImageAPIMethod.swift +++ b/FirebaseAI/Sources/ImageAPIMethod.swift @@ -16,5 +16,5 @@ import Foundation enum ImageAPIMethod: String { - case generateImages = "images:generate" + case generateImages = "templatePredict" } diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift index 86fdcfcb3cf..46eb73013b8 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift @@ -43,11 +43,11 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { } func testGenerateImages() async throws { - let imagenModel = FirebaseAI.firebaseAI(backend: .googleAI()).templateImagenModel() + let imagenModel = FirebaseAI.firebaseAI(backend: .vertexAI()).templateImagenModel() let imagenPrompt = "A cat picture" do { let response = try await imagenModel.generateImages( - template: "describeImages", + template: "generate_images", variables: [ "prompt": imagenPrompt, ] From d2018a69fc8cd0417b04d3610b983684b1d554e1 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Sat, 11 Oct 2025 20:29:00 -0700 Subject: [PATCH 13/24] Imagen generate test works --- .../Sources/GenerateImagesRequest.swift | 2 +- FirebaseAI/Sources/TemplateImagenModel.swift | 25 ++++--------------- ...ServerPromptTemplateIntegrationTests.swift | 2 +- 3 files changed, 7 insertions(+), 22 deletions(-) diff --git a/FirebaseAI/Sources/GenerateImagesRequest.swift b/FirebaseAI/Sources/GenerateImagesRequest.swift index 4343589bdb2..80f2dbbf8a5 100644 --- a/FirebaseAI/Sources/GenerateImagesRequest.swift +++ b/FirebaseAI/Sources/GenerateImagesRequest.swift @@ -17,7 +17,7 @@ import Foundation @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) public class GenerateImagesRequest: @unchecked Sendable, GenerativeAIRequest { - public typealias Response = GenerateImagesResponse + public typealias Response = ImagenGenerationResponse public var url: URL { var urlString = diff --git a/FirebaseAI/Sources/TemplateImagenModel.swift b/FirebaseAI/Sources/TemplateImagenModel.swift index 3c8efc31425..845b0ba2e10 100644 --- a/FirebaseAI/Sources/TemplateImagenModel.swift +++ b/FirebaseAI/Sources/TemplateImagenModel.swift @@ -1,4 +1,3 @@ - // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); @@ -48,31 +47,17 @@ public final class TemplateImagenModel: Sendable { apiConfig: apiConfig, options: options ) - let response: GenerateImagesResponse = try await generativeAIService + let response: ImagenGenerationResponse = try await generativeAIService .loadRequest(request: request) - return response + return GenerateImagesResponse(images: response.images.map(\.data)) } } // A placeholder for the response from an image generation request. -public struct GenerateImagesResponse: Decodable, @unchecked Sendable { +public struct GenerateImagesResponse: @unchecked Sendable { public let images: [Data] - enum CodingKeys: String, CodingKey { - case predictions - } - - struct Prediction: Decodable { - let images: [String] - } - - public init(from decoder: Decoder) throws { - let container = try decoder.container(keyedBy: CodingKeys.self) - let predictions = try container.decode([Prediction].self, forKey: .predictions) - guard let firstPrediction = predictions.first else { - images = [] - return - } - images = firstPrediction.images.compactMap { Data(base64Encoded: $0) } + public init(images: [Data]) { + self.images = images } } diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift index 46eb73013b8..61d804926d5 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift @@ -52,7 +52,7 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { "prompt": imagenPrompt, ] ) - XCTAssertEqual(response.images.count, 1) + XCTAssertEqual(response.images.count, 3) } catch { XCTFail("An error occurred: \(error)") } From 7a8b7adc9d1acd20606f62cf522758ee0c97c166 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Sun, 12 Oct 2025 10:21:31 -0700 Subject: [PATCH 14/24] fix template imagen response return type --- FirebaseAI/Sources/TemplateImagenModel.swift | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/FirebaseAI/Sources/TemplateImagenModel.swift b/FirebaseAI/Sources/TemplateImagenModel.swift index 845b0ba2e10..10d09eaa18c 100644 --- a/FirebaseAI/Sources/TemplateImagenModel.swift +++ b/FirebaseAI/Sources/TemplateImagenModel.swift @@ -37,7 +37,7 @@ public final class TemplateImagenModel: Sendable { public func generateImages(template: String, variables: [String: Any], options: RequestOptions = RequestOptions()) async throws - -> GenerateImagesResponse { + -> ImagenGenerationResponse { let templateVariables = try variables.mapValues { try TemplateVariable(value: $0) } let projectID = generativeAIService.firebaseInfo.projectID let request = GenerateImagesRequest( @@ -47,17 +47,6 @@ public final class TemplateImagenModel: Sendable { apiConfig: apiConfig, options: options ) - let response: ImagenGenerationResponse = try await generativeAIService - .loadRequest(request: request) - return GenerateImagesResponse(images: response.images.map(\.data)) + return try await generativeAIService.loadRequest(request: request) } -} - -// A placeholder for the response from an image generation request. -public struct GenerateImagesResponse: @unchecked Sendable { - public let images: [Data] - - public init(images: [Data]) { - self.images = images - } -} +} \ No newline at end of file From 74a8c2d998dfa2af356e333f14add21cd8a7da1d Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Sun, 12 Oct 2025 10:49:56 -0700 Subject: [PATCH 15/24] existing tests passing --- FirebaseAI/Sources/History.swift | 3 +-- FirebaseAI/Sources/ImageAPIMethod.swift | 2 -- FirebaseAI/Sources/TemplateVariable.swift | 1 - .../ServerPromptTemplateIntegrationTests.swift | 8 ++++---- 4 files changed, 5 insertions(+), 9 deletions(-) diff --git a/FirebaseAI/Sources/History.swift b/FirebaseAI/Sources/History.swift index 7e5b81ad60b..827f7df5b46 100644 --- a/FirebaseAI/Sources/History.swift +++ b/FirebaseAI/Sources/History.swift @@ -1,5 +1,4 @@ - -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. diff --git a/FirebaseAI/Sources/ImageAPIMethod.swift b/FirebaseAI/Sources/ImageAPIMethod.swift index 7a5efbaa032..bf8a6c367d1 100644 --- a/FirebaseAI/Sources/ImageAPIMethod.swift +++ b/FirebaseAI/Sources/ImageAPIMethod.swift @@ -1,9 +1,7 @@ - // Copyright 2025 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 not use this file except in compliance with the License. // // http://www.apache.org/licenses/LICENSE-2.0 // diff --git a/FirebaseAI/Sources/TemplateVariable.swift b/FirebaseAI/Sources/TemplateVariable.swift index 7ccdc381f83..e89147e6ace 100644 --- a/FirebaseAI/Sources/TemplateVariable.swift +++ b/FirebaseAI/Sources/TemplateVariable.swift @@ -1,4 +1,3 @@ - // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift index 61d804926d5..8f61c628fad 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift @@ -1,4 +1,3 @@ - // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); @@ -59,7 +58,7 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { } func testGenerateContentWithMedia() async throws { - let model = FirebaseAI.firebaseAI(backend: .googleAI()).templateGenerativeModel() + let model = FirebaseAI.firebaseAI(backend: .vertexAI()).templateGenerativeModel() let image = UIImage(systemName: "photo")! if let imageBytes = image.jpegData(compressionQuality: 0.8) { let base64Image = imageBytes.base64EncodedString() @@ -85,7 +84,7 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { } func testChat() async throws { - let model = FirebaseAI.firebaseAI(backend: .googleAI()).templateGenerativeModel() + let model = FirebaseAI.firebaseAI(backend: .vertexAI()).templateGenerativeModel() let initialHistory = [ ModelContent(role: "user", parts: "Hello!"), ModelContent(role: "model", parts: "Hi there! How can I help?"), @@ -100,7 +99,8 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { variables: ["message": userMessage] ) XCTAssert(response.text?.isEmpty == false) - XCTAssertEqual(chatSession.history.count, 3) + XCTAssertEqual(chatSession.history.count, 4) + XCTAssertEqual((chatSession.history[2].parts.first as? TextPart)?.text, userMessage) } catch { XCTFail("An error occurred: \(error)") } From 09ca09aba1837347c41ea2e6c1a6e839815e43e4 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Sun, 12 Oct 2025 15:53:50 -0700 Subject: [PATCH 16/24] build fix --- FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift index 1456bd4cea8..2deddcb95ff 100644 --- a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift @@ -48,6 +48,6 @@ final class TemplateImagenModelTests: XCTestCase { variables: ["prompt": "a cat picture"] ) XCTAssertEqual(response.images.count, 1) - XCTAssertEqual(response.images.first, Data(base64Encoded: "aW1hZ2UgZGF0YQ==")) + XCTAssertEqual(response.images.first?.data, Data(base64Encoded: "aW1hZ2UgZGF0YQ==")) } } From 5ad9314118397242fecba4ba9a2b87e5edb712f8 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Sun, 12 Oct 2025 15:58:43 -0700 Subject: [PATCH 17/24] style --- FirebaseAI/Sources/TemplateImagenModel.swift | 2 +- FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift | 1 - FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift | 1 - FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift | 1 - 4 files changed, 1 insertion(+), 4 deletions(-) diff --git a/FirebaseAI/Sources/TemplateImagenModel.swift b/FirebaseAI/Sources/TemplateImagenModel.swift index 10d09eaa18c..0e37d1689b1 100644 --- a/FirebaseAI/Sources/TemplateImagenModel.swift +++ b/FirebaseAI/Sources/TemplateImagenModel.swift @@ -49,4 +49,4 @@ public final class TemplateImagenModel: Sendable { ) return try await generativeAIService.loadRequest(request: request) } -} \ No newline at end of file +} diff --git a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift index 4bd920c0caf..f8bdc8a8ed8 100644 --- a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift @@ -1,4 +1,3 @@ - // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift b/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift index c5928fb7aa6..6371c742430 100644 --- a/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift @@ -1,4 +1,3 @@ - // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift index 2deddcb95ff..c5c08ed1629 100644 --- a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift @@ -1,4 +1,3 @@ - // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); From a0ac6c52d0ec53a22b8a5ee835c1b65be43cc405 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Mon, 13 Oct 2025 14:31:31 -0700 Subject: [PATCH 18/24] Add streaming integration tests --- FirebaseAI/Sources/APIMethod.swift | 1 - ...ServerPromptTemplateIntegrationTests.swift | 57 ++++++++++++++++++- 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/FirebaseAI/Sources/APIMethod.swift b/FirebaseAI/Sources/APIMethod.swift index 9ae05e88429..9afa9c163aa 100644 --- a/FirebaseAI/Sources/APIMethod.swift +++ b/FirebaseAI/Sources/APIMethod.swift @@ -1,4 +1,3 @@ - // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift index 8f61c628fad..e22301c946e 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift @@ -30,7 +30,7 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { template: "greeting", variables: [ "name": userName, - "language": "English", + "language": "Spanish", ] ) let text = try XCTUnwrap(response.text) @@ -41,6 +41,29 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { } } + func testGenerateContentStream() async throws { + let model = FirebaseAI.firebaseAI(backend: .vertexAI()).templateGenerativeModel() + let userName = "paul" + do { + let stream = try model.generateContentStream( + template: "greeting", + variables: [ + "name": userName, + "language": "English", + ] + ) + var resultText = "" + for try await response in stream { + if let text = response.text { + resultText += text + } + } + XCTAssert(resultText.contains("Paul")) + } catch { + XCTFail("An error occurred: \(error)") + } + } + func testGenerateImages() async throws { let imagenModel = FirebaseAI.firebaseAI(backend: .vertexAI()).templateImagenModel() let imagenPrompt = "A cat picture" @@ -83,6 +106,38 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { } } + func testGenerateContentStreamWithMedia() async throws { + let model = FirebaseAI.firebaseAI(backend: .vertexAI()).templateGenerativeModel() + let image = UIImage(systemName: "photo")! + if let imageBytes = image.jpegData(compressionQuality: 0.8) { + let base64Image = imageBytes.base64EncodedString() + + do { + let stream = try model.generateContentStream( + template: "media", + variables: [ + "imageData": [ + "isInline": true, + "mimeType": "image/jpeg", + "contents": base64Image, + ], + ] + ) + var resultText = "" + for try await response in stream { + if let text = response.text { + resultText += text + } + } + XCTAssert(resultText.isEmpty == false) + } catch { + XCTFail("An error occurred: \(error)") + } + } else { + XCTFail("Could not get image data.") + } + } + func testChat() async throws { let model = FirebaseAI.firebaseAI(backend: .vertexAI()).templateGenerativeModel() let initialHistory = [ From ee1818ec9c36f6987b4ff3638deb52e25af51785 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Mon, 13 Oct 2025 14:39:27 -0700 Subject: [PATCH 19/24] testChatStream --- FirebaseAI/Sources/TemplateChatSession.swift | 2 +- .../Sources/TemplateGenerativeModel.swift | 1 - ...ServerPromptTemplateIntegrationTests.swift | 29 +++++++++++++++++++ 3 files changed, 30 insertions(+), 2 deletions(-) diff --git a/FirebaseAI/Sources/TemplateChatSession.swift b/FirebaseAI/Sources/TemplateChatSession.swift index fabd4fcd52b..e6e705f5814 100644 --- a/FirebaseAI/Sources/TemplateChatSession.swift +++ b/FirebaseAI/Sources/TemplateChatSession.swift @@ -1,4 +1,4 @@ -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. diff --git a/FirebaseAI/Sources/TemplateGenerativeModel.swift b/FirebaseAI/Sources/TemplateGenerativeModel.swift index 98df95fba9c..24b6d638eee 100644 --- a/FirebaseAI/Sources/TemplateGenerativeModel.swift +++ b/FirebaseAI/Sources/TemplateGenerativeModel.swift @@ -1,4 +1,3 @@ - // Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift index e22301c946e..e3a04424d99 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift @@ -160,4 +160,33 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { XCTFail("An error occurred: \(error)") } } + + func testChatStream() async throws { + let model = FirebaseAI.firebaseAI(backend: .vertexAI()).templateGenerativeModel() + let initialHistory = [ + ModelContent(role: "user", parts: "Hello!"), + ModelContent(role: "model", parts: "Hi there! How can I help?"), + ] + let chatSession = model.startChat(template: "chat_history", history: initialHistory) + + let userMessage = "What's the weather like?" + + do { + let stream = try chatSession.sendMessageStream( + userMessage, + variables: ["message": userMessage] + ) + var resultText = "" + for try await response in stream { + if let text = response.text { + resultText += text + } + } + XCTAssert(resultText.isEmpty == false) + XCTAssertEqual(chatSession.history.count, 4) + XCTAssertEqual((chatSession.history[2].parts.first as? TextPart)?.text, userMessage) + } catch { + XCTFail("An error occurred: \(error)") + } + } } From d36d33ded7c481afa7a3dd5decd947382661c284 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Mon, 13 Oct 2025 14:49:32 -0700 Subject: [PATCH 20/24] unit tests run --- .../Tests/Unit/TemplateChatSessionTests.swift | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift index f8bdc8a8ed8..281a0cabf18 100644 --- a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift @@ -19,19 +19,29 @@ import XCTest @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) final class TemplateChatSessionTests: XCTestCase { var model: TemplateGenerativeModel! + var urlSession: URLSession! override func setUp() { super.setUp() + let configuration = URLSessionConfiguration.default + configuration.protocolClasses = [MockURLProtocol.self] + urlSession = URLSession(configuration: configuration) let firebaseInfo = GenerativeModelTestUtil.testFirebaseInfo() let generativeAIService = GenerativeAIService( firebaseInfo: firebaseInfo, - urlSession: GenAIURLSession.default + urlSession: urlSession ) let apiConfig = APIConfig(service: .googleAI(endpoint: .firebaseProxyProd), version: .v1beta) model = TemplateGenerativeModel(generativeAIService: generativeAIService, apiConfig: apiConfig) } func testSendMessage() async throws { + MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( + forResource: "unary-success-response", + withExtension: "json", + subdirectory: "mock-responses", + isTemplateRequest: true + ) let chat = model.startChat(template: "test-template") let response = try await chat.sendMessage("Hello", variables: ["name": "test"]) XCTAssertEqual(chat.history.count, 2) @@ -40,4 +50,4 @@ final class TemplateChatSessionTests: XCTestCase { XCTAssertEqual(chat.history[1].role, "model") XCTAssertEqual(response.candidates.count, 1) } -} +} \ No newline at end of file From 9d7fc0c9dcbe42353cc67c97cbb6c8db984638b7 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Mon, 13 Oct 2025 14:57:59 -0700 Subject: [PATCH 21/24] streaming unit tests --- .../Tests/Unit/TemplateChatSessionTests.swift | 26 ++++++++++++++++++- .../Unit/TemplateGenerativeModelTests.swift | 22 ++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift index 281a0cabf18..2dc9c071c3c 100644 --- a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift @@ -50,4 +50,28 @@ final class TemplateChatSessionTests: XCTestCase { XCTAssertEqual(chat.history[1].role, "model") XCTAssertEqual(response.candidates.count, 1) } -} \ No newline at end of file + + func testSendMessageStream() async throws { + MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( + forResource: "streaming-success-basic-reply-short", + withExtension: "txt", + subdirectory: "mock-responses/googleai", + isTemplateRequest: true + ) + let chat = model.startChat(template: "test-template") + let stream = try chat.sendMessageStream("Hello", variables: ["name": "test"]) + + var content = "" + for try await response in stream { + if let text = response.text { + content += text + } + } + + XCTAssertEqual(content, "The capital of Wyoming is **Cheyenne**.\n") + XCTAssertEqual(chat.history.count, 2) + XCTAssertEqual(chat.history[0].role, "user") + XCTAssertEqual((chat.history[0].parts.first as? TextPart)?.text, "Hello") + XCTAssertEqual(chat.history[1].role, "model") + } +} diff --git a/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift b/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift index 6371c742430..0d42851d6f5 100644 --- a/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift @@ -49,4 +49,26 @@ final class TemplateGenerativeModelTests: XCTestCase { ) XCTAssertEqual(response.text, "Hello, world!") } + + func testGenerateContentStream() async throws { + MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( + forResource: "streaming-success-basic-reply-short", + withExtension: "txt", + subdirectory: "mock-responses/googleai", + isTemplateRequest: true + ) + + let stream = try model.generateContentStream( + template: "test-template", + variables: ["name": "test"] + ) + + var content = "" + for try await response in stream { + if let text = response.text { + content += text + } + } + XCTAssertEqual(content, "The capital of Wyoming is **Cheyenne**.\n") + } } From 812c05aef2113b7bd053b311cd36515b0e6876a8 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Mon, 13 Oct 2025 15:37:23 -0700 Subject: [PATCH 22/24] existing mock files for template unit tests --- FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift | 8 ++++++-- FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift | 9 ++++++--- FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift | 8 ++++---- 3 files changed, 16 insertions(+), 9 deletions(-) diff --git a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift index 2dc9c071c3c..33481c0dab6 100644 --- a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift @@ -37,9 +37,9 @@ final class TemplateChatSessionTests: XCTestCase { func testSendMessage() async throws { MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( - forResource: "unary-success-response", + forResource: "unary-success-basic-reply-short", withExtension: "json", - subdirectory: "mock-responses", + subdirectory: "mock-responses/googleai", isTemplateRequest: true ) let chat = model.startChat(template: "test-template") @@ -48,6 +48,10 @@ final class TemplateChatSessionTests: XCTestCase { XCTAssertEqual(chat.history[0].role, "user") XCTAssertEqual((chat.history[0].parts.first as? TextPart)?.text, "Hello") XCTAssertEqual(chat.history[1].role, "model") + XCTAssertEqual( + (chat.history[1].parts.first as? TextPart)?.text, + "Google's headquarters, also known as the Googleplex, is located in **Mountain View, California**.\n" + ) XCTAssertEqual(response.candidates.count, 1) } diff --git a/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift b/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift index 0d42851d6f5..841bfd7763e 100644 --- a/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift @@ -37,9 +37,9 @@ final class TemplateGenerativeModelTests: XCTestCase { func testGenerateContent() async throws { MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( - forResource: "unary-success-response", + forResource: "unary-success-basic-reply-short", withExtension: "json", - subdirectory: "mock-responses", + subdirectory: "mock-responses/googleai", isTemplateRequest: true ) @@ -47,7 +47,10 @@ final class TemplateGenerativeModelTests: XCTestCase { template: "test-template", variables: ["name": "test"] ) - XCTAssertEqual(response.text, "Hello, world!") + XCTAssertEqual( + response.text, + "Google's headquarters, also known as the Googleplex, is located in **Mountain View, California**.\n" + ) } func testGenerateContentStream() async throws { diff --git a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift index c5c08ed1629..f350d4ae980 100644 --- a/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift @@ -36,9 +36,9 @@ final class TemplateImagenModelTests: XCTestCase { func testGenerateImages() async throws { MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( - forResource: "unary-success-image-response", + forResource: "unary-success-generate-images-base64", withExtension: "json", - subdirectory: "mock-responses", + subdirectory: "mock-responses/vertexai", isTemplateRequest: true ) @@ -46,7 +46,7 @@ final class TemplateImagenModelTests: XCTestCase { template: "test-template", variables: ["prompt": "a cat picture"] ) - XCTAssertEqual(response.images.count, 1) - XCTAssertEqual(response.images.first?.data, Data(base64Encoded: "aW1hZ2UgZGF0YQ==")) + XCTAssertEqual(response.images.count, 4) + XCTAssertNotNil(response.images.first?.data) } } From df1f4d4a0c27e7563fd9feca717340468d1bed0c Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Mon, 13 Oct 2025 16:08:12 -0700 Subject: [PATCH 23/24] streamline integration tests --- .../TemplateGenerateContentRequest.swift | 2 - ...ServerPromptTemplateIntegrationTests.swift | 176 ++++++++---------- 2 files changed, 74 insertions(+), 104 deletions(-) diff --git a/FirebaseAI/Sources/TemplateGenerateContentRequest.swift b/FirebaseAI/Sources/TemplateGenerateContentRequest.swift index 7c2da089b80..65e4708e96b 100644 --- a/FirebaseAI/Sources/TemplateGenerateContentRequest.swift +++ b/FirebaseAI/Sources/TemplateGenerateContentRequest.swift @@ -21,7 +21,6 @@ struct TemplateGenerateContentRequest: Sendable { let history: [ModelContent] let projectID: String let stream: Bool - let apiConfig: APIConfig let options: RequestOptions } @@ -53,7 +52,6 @@ extension TemplateGenerateContentRequest: GenerativeAIRequest { let templateName = template.hasSuffix(".prompt") ? template : "\(template).prompt" urlString += "/templates/\(templateName):templateGenerateContent" if stream { - // TODO: Fix this. urlString += "?alt=sse" } return URL(string: urlString)! diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift index e3a04424d99..84072b65980 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift @@ -25,59 +25,47 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { func testGenerateContentWithText() async throws { let model = FirebaseAI.firebaseAI(backend: .vertexAI()).templateGenerativeModel() let userName = "paul" - do { - let response = try await model.generateContent( - template: "greeting", - variables: [ - "name": userName, - "language": "Spanish", - ] - ) - let text = try XCTUnwrap(response.text) - print(text) - XCTAssert(text.contains("Paul")) - } catch { - XCTFail("An error occurred: \(error)") - } + let response = try await model.generateContent( + template: "greeting", + variables: [ + "name": userName, + "language": "Spanish", + ] + ) + let text = try XCTUnwrap(response.text) + print(text) + XCTAssert(text.contains("Paul")) } func testGenerateContentStream() async throws { let model = FirebaseAI.firebaseAI(backend: .vertexAI()).templateGenerativeModel() let userName = "paul" - do { - let stream = try model.generateContentStream( - template: "greeting", - variables: [ - "name": userName, - "language": "English", - ] - ) - var resultText = "" - for try await response in stream { - if let text = response.text { - resultText += text - } + let stream = try model.generateContentStream( + template: "greeting", + variables: [ + "name": userName, + "language": "English", + ] + ) + var resultText = "" + for try await response in stream { + if let text = response.text { + resultText += text } - XCTAssert(resultText.contains("Paul")) - } catch { - XCTFail("An error occurred: \(error)") } + XCTAssert(resultText.contains("Paul")) } func testGenerateImages() async throws { let imagenModel = FirebaseAI.firebaseAI(backend: .vertexAI()).templateImagenModel() let imagenPrompt = "A cat picture" - do { - let response = try await imagenModel.generateImages( - template: "generate_images", - variables: [ - "prompt": imagenPrompt, - ] - ) - XCTAssertEqual(response.images.count, 3) - } catch { - XCTFail("An error occurred: \(error)") - } + let response = try await imagenModel.generateImages( + template: "generate_images", + variables: [ + "prompt": imagenPrompt, + ] + ) + XCTAssertEqual(response.images.count, 3) } func testGenerateContentWithMedia() async throws { @@ -86,21 +74,17 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { if let imageBytes = image.jpegData(compressionQuality: 0.8) { let base64Image = imageBytes.base64EncodedString() - do { - let response = try await model.generateContent( - template: "media", - variables: [ - "imageData": [ - "isInline": true, - "mimeType": "image/jpeg", - "contents": base64Image, - ], - ] - ) - XCTAssert(response.text?.isEmpty == false) - } catch { - XCTFail("An error occurred: \(error)") - } + let response = try await model.generateContent( + template: "media", + variables: [ + "imageData": [ + "isInline": true, + "mimeType": "image/jpeg", + "contents": base64Image, + ], + ] + ) + XCTAssert(response.text?.isEmpty == false) } else { XCTFail("Could not get image data.") } @@ -112,27 +96,23 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { if let imageBytes = image.jpegData(compressionQuality: 0.8) { let base64Image = imageBytes.base64EncodedString() - do { - let stream = try model.generateContentStream( - template: "media", - variables: [ - "imageData": [ - "isInline": true, - "mimeType": "image/jpeg", - "contents": base64Image, - ], - ] - ) - var resultText = "" - for try await response in stream { - if let text = response.text { - resultText += text - } + let stream = try model.generateContentStream( + template: "media", + variables: [ + "imageData": [ + "isInline": true, + "mimeType": "image/jpeg", + "contents": base64Image, + ], + ] + ) + var resultText = "" + for try await response in stream { + if let text = response.text { + resultText += text } - XCTAssert(resultText.isEmpty == false) - } catch { - XCTFail("An error occurred: \(error)") } + XCTAssert(resultText.isEmpty == false) } else { XCTFail("Could not get image data.") } @@ -148,17 +128,13 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { let userMessage = "What's the weather like?" - do { - let response = try await chatSession.sendMessage( - userMessage, - variables: ["message": userMessage] - ) - XCTAssert(response.text?.isEmpty == false) - XCTAssertEqual(chatSession.history.count, 4) - XCTAssertEqual((chatSession.history[2].parts.first as? TextPart)?.text, userMessage) - } catch { - XCTFail("An error occurred: \(error)") - } + let response = try await chatSession.sendMessage( + userMessage, + variables: ["message": userMessage] + ) + XCTAssert(response.text?.isEmpty == false) + XCTAssertEqual(chatSession.history.count, 4) + XCTAssertEqual((chatSession.history[2].parts.first as? TextPart)?.text, userMessage) } func testChatStream() async throws { @@ -171,22 +147,18 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { let userMessage = "What's the weather like?" - do { - let stream = try chatSession.sendMessageStream( - userMessage, - variables: ["message": userMessage] - ) - var resultText = "" - for try await response in stream { - if let text = response.text { - resultText += text - } + let stream = try chatSession.sendMessageStream( + userMessage, + variables: ["message": userMessage] + ) + var resultText = "" + for try await response in stream { + if let text = response.text { + resultText += text } - XCTAssert(resultText.isEmpty == false) - XCTAssertEqual(chatSession.history.count, 4) - XCTAssertEqual((chatSession.history[2].parts.first as? TextPart)?.text, userMessage) - } catch { - XCTFail("An error occurred: \(error)") } + XCTAssert(resultText.isEmpty == false) + XCTAssertEqual(chatSession.history.count, 4) + XCTAssertEqual((chatSession.history[2].parts.first as? TextPart)?.text, userMessage) } -} +} \ No newline at end of file From ad2a87bbb3b1c95087620e123cbb2faa42c641bd Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Mon, 13 Oct 2025 16:58:24 -0700 Subject: [PATCH 24/24] fixes --- FirebaseAI/Sources/TemplateVariable.swift | 2 ++ .../ServerPromptTemplateIntegrationTests.swift | 3 +-- .../Tests/Unit/TemplateChatSessionTests.swift | 7 +------ .../Tests/Unit/TemplateGenerativeModelTests.swift | 7 +------ .../TestUtilities/GenerativeModelTestUtil.swift | 13 +++++++++++++ 5 files changed, 18 insertions(+), 14 deletions(-) diff --git a/FirebaseAI/Sources/TemplateVariable.swift b/FirebaseAI/Sources/TemplateVariable.swift index e89147e6ace..a942c2c715b 100644 --- a/FirebaseAI/Sources/TemplateVariable.swift +++ b/FirebaseAI/Sources/TemplateVariable.swift @@ -30,6 +30,8 @@ enum TemplateVariable: Encodable, Sendable { self = .int(value) case let value as Double: self = .double(value) + case let value as Float: + self = .double(Double(value)) case let value as Bool: self = .bool(value) case let value as [Any]: diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift index 84072b65980..69b6d1add3a 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift @@ -33,7 +33,6 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { ] ) let text = try XCTUnwrap(response.text) - print(text) XCTAssert(text.contains("Paul")) } @@ -161,4 +160,4 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase { XCTAssertEqual(chatSession.history.count, 4) XCTAssertEqual((chatSession.history[2].parts.first as? TextPart)?.text, userMessage) } -} \ No newline at end of file +} diff --git a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift index 33481c0dab6..28278549f4d 100644 --- a/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift @@ -65,12 +65,7 @@ final class TemplateChatSessionTests: XCTestCase { let chat = model.startChat(template: "test-template") let stream = try chat.sendMessageStream("Hello", variables: ["name": "test"]) - var content = "" - for try await response in stream { - if let text = response.text { - content += text - } - } + let content = try await GenerativeModelTestUtil.collectTextFromStream(stream) XCTAssertEqual(content, "The capital of Wyoming is **Cheyenne**.\n") XCTAssertEqual(chat.history.count, 2) diff --git a/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift b/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift index 841bfd7763e..b191ccf54cb 100644 --- a/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift +++ b/FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift @@ -66,12 +66,7 @@ final class TemplateGenerativeModelTests: XCTestCase { variables: ["name": "test"] ) - var content = "" - for try await response in stream { - if let text = response.text { - content += text - } - } + let content = try await GenerativeModelTestUtil.collectTextFromStream(stream) XCTAssertEqual(content, "The capital of Wyoming is **Cheyenne**.\n") } } diff --git a/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift b/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift index e783fdd6388..8ef28d33114 100644 --- a/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift +++ b/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift @@ -88,6 +88,19 @@ enum GenerativeModelTestUtil { #endif // os(watchOS) } + static func collectTextFromStream(_ stream: AsyncThrowingStream< + GenerateContentResponse, + Error + >) async throws -> String { + var content = "" + for try await response in stream { + if let text = response.text { + content += text + } + } + return content + } + static func nonHTTPRequestHandler() throws -> ((URLRequest) -> ( URLResponse, AsyncLineSequence?