[Vertex AI] Test SDK with v1 API instead of v1beta

firebase · Jan 14, 2025 · 8ef93d3 · 8ef93d3
1 parent b756d39
commit 8ef93d3
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git 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
@@ -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() }