From 8ef93d3a04bfc81e86a202bb7260e421d7c3641a Mon Sep 17 00:00:00 2001 From: Andrew Heard Date: Tue, 14 Jan 2025 11:00:23 -0500 Subject: [PATCH 1/2] [Vertex AI] Test SDK with `v1` API instead of `v1beta` --- FirebaseVertexAI/Sources/GenerativeAIRequest.swift | 2 +- .../Tests/Integration/IntegrationTests.swift | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/FirebaseVertexAI/Sources/GenerativeAIRequest.swift b/FirebaseVertexAI/Sources/GenerativeAIRequest.swift index b792830120e..f0c7cecde8e 100644 --- a/FirebaseVertexAI/Sources/GenerativeAIRequest.swift +++ b/FirebaseVertexAI/Sources/GenerativeAIRequest.swift @@ -31,7 +31,7 @@ public struct RequestOptions { let timeout: TimeInterval /// The API version to use in requests to the backend. - let apiVersion = "v1beta" + let apiVersion = "v1" /// Initializes a request options object. /// diff --git a/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift b/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift index 8e6e6c8d601..a09d6a99e79 100644 --- a/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift +++ b/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift @@ -75,6 +75,19 @@ final class IntegrationTests: XCTestCase { XCTAssertEqual(text, "Mountain View") } + func testGenerateContentStream() async throws { + let prompt = "Where is Google headquarters located? Answer with the city name only." + + var text = "" + let contentStream = try model.generateContentStream(prompt) + for try await chunk in contentStream { + text += try XCTUnwrap(chunk.text) + } + + text = text.trimmingCharacters(in: .whitespacesAndNewlines) + XCTAssertEqual(text, "Mountain View") + } + func testGenerateContent_appCheckNotConfigured_shouldFail() async throws { let app = try FirebaseApp.defaultNamedCopy(name: TestAppCheckProviderFactory.notConfiguredName) addTeardownBlock { await app.delete() } From c542d621206a1e7562e38877ffe4af933ff6e7d0 Mon Sep 17 00:00:00 2001 From: Andrew Heard Date: Tue, 14 Jan 2025 11:40:47 -0500 Subject: [PATCH 2/2] Add more integration tests --- .../Tests/Integration/IntegrationTests.swift | 94 +++++++++++++++---- 1 file changed, 78 insertions(+), 16 deletions(-) diff --git a/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift b/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift index a09d6a99e79..53d2fdb192d 100644 --- a/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift +++ b/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift @@ -66,7 +66,7 @@ final class IntegrationTests: XCTestCase { // MARK: - Generate Content - func testGenerateContent() async throws { + func testGenerateContent_text() async throws { let prompt = "Where is Google headquarters located? Answer with the city name only." let response = try await model.generateContent(prompt) @@ -75,17 +75,29 @@ final class IntegrationTests: XCTestCase { XCTAssertEqual(text, "Mountain View") } - func testGenerateContentStream() async throws { - let prompt = "Where is Google headquarters located? Answer with the city name only." + func testGenerateContent_image_fileData_public() async throws { + let storageRef = storage.reference(withPath: "vertexai/public/green.png") + let fileData = FileDataPart(uri: storageRef.gsURI, mimeType: "image/png") - var text = "" - let contentStream = try model.generateContentStream(prompt) - for try await chunk in contentStream { - text += try XCTUnwrap(chunk.text) - } + let response = try await model.generateContent(fileData, "What color is this?") - text = text.trimmingCharacters(in: .whitespacesAndNewlines) - XCTAssertEqual(text, "Mountain View") + XCTAssertNotNil(response.text) + } + + func testGenerateContent_image_fileData_requiresUserAuth_wrongUser_permissionDenied() async throws { + let userID = "3MjEzU6JIobWvHdCYHicnDMcPpQ2" + let storageRef = storage.reference(withPath: "vertexai/authenticated/user/\(userID)/pink.webp") + + let fileData = FileDataPart(uri: storageRef.gsURI, mimeType: "image/webp") + + do { + let response = try await model.generateContent(fileData, "What color is this?") + XCTFail("Expected to throw an error, got response: \(response)") + } catch { + let errorDescription = String(describing: error) + XCTAssertTrue(errorDescription.contains("403")) + XCTAssertTrue(errorDescription.contains("The caller does not have permission")) + } } func testGenerateContent_appCheckNotConfigured_shouldFail() async throws { @@ -96,13 +108,63 @@ final class IntegrationTests: XCTestCase { let prompt = "Where is Google headquarters located? Answer with the city name only." do { - _ = try await model.generateContent(prompt) - XCTFail("Expected a Firebase App Check error; none thrown.") + let response = try await model.generateContent(prompt) + XCTFail("Expected a Firebase App Check error, got response: \(response)") } catch let GenerateContentError.internalError(error) { XCTAssertTrue(String(describing: error).contains("Firebase App Check token is invalid")) } } + // MARK: - Generate Content Streaming + + func testGenerateContentStream_text() async throws { + let prompt = "Where is Google headquarters located? Answer with the city name only." + + var text = "" + let contentStream = try model.generateContentStream(prompt) + for try await chunk in contentStream { + text += try XCTUnwrap(chunk.text) + } + + text = text.trimmingCharacters(in: .whitespacesAndNewlines) + XCTAssertEqual(text, "Mountain View") + } + + func testGenerateContentStream_image_fileData_public() async throws { + let storageRef = storage.reference(withPath: "vertexai/public/green.png") + let fileData = FileDataPart(uri: storageRef.gsURI, mimeType: "image/png") + + var text = "" + let contentStream = try model.generateContentStream(fileData, "What color is this?") + for try await chunk in contentStream { + text += try XCTUnwrap(chunk.text) + } + + text = text.trimmingCharacters(in: .whitespacesAndNewlines) + XCTAssertFalse(text.isEmpty) + } + + func testGenerateContentStream_image_fileData_requiresUserAuth_wrongUser_permissionDenied() async throws { + let userID = "3MjEzU6JIobWvHdCYHicnDMcPpQ2" + let storageRef = storage.reference(withPath: "vertexai/authenticated/user/\(userID)/pink.webp") + + let fileData = FileDataPart(uri: storageRef.gsURI, mimeType: "image/webp") + + let contentStream = try model.generateContentStream(fileData, "What color is this?") + + do { + var responses = [GenerateContentResponse]() + for try await response in contentStream { + responses.append(response) + } + XCTFail("Expected to throw an error, got response(s): \(responses)") + } catch { + let errorDescription = String(describing: error) + XCTAssertTrue(errorDescription.contains("403")) + XCTAssertTrue(errorDescription.contains("The caller does not have permission")) + } + } + // MARK: - Count Tokens func testCountTokens_text() async throws { @@ -179,8 +241,8 @@ final class IntegrationTests: XCTestCase { let fileData = FileDataPart(uri: storageRef.gsURI, mimeType: "image/webp") do { - _ = try await model.countTokens(fileData) - XCTFail("Expected to throw an error.") + let response = try await model.countTokens(fileData) + XCTFail("Expected to throw an error, got response: \(response)") } catch { let errorDescription = String(describing: error) XCTAssertTrue(errorDescription.contains("403")) @@ -242,8 +304,8 @@ final class IntegrationTests: XCTestCase { let prompt = "Why is the sky blue?" do { - _ = try await model.countTokens(prompt) - XCTFail("Expected a Firebase App Check error; none thrown.") + let response = try await model.countTokens(prompt) + XCTFail("Expected a Firebase App Check error, got response: \(response)") } catch { XCTAssertTrue(String(describing: error).contains("Firebase App Check token is invalid")) }