Skip to content

Commit 996cd4b

Browse files
committed
Add an integration test and fix uncovered camelCase JSON issues
1 parent ee3fde0 commit 996cd4b

File tree

3 files changed

+47
-1
lines changed

3 files changed

+47
-1
lines changed

FirebaseVertexAI/Sources/GenerationConfig.swift

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -162,4 +162,17 @@ public struct GenerationConfig {
162162
// MARK: - Codable Conformances
163163

164164
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
165-
extension GenerationConfig: Encodable {}
165+
extension GenerationConfig: Encodable {
166+
enum CodingKeys: String, CodingKey {
167+
case temperature
168+
case topP
169+
case topK
170+
case candidateCount
171+
case maxOutputTokens
172+
case presencePenalty
173+
case frequencyPenalty
174+
case stopSequences
175+
case responseMIMEType = "responseMimeType"
176+
case responseSchema
177+
}
178+
}

FirebaseVertexAI/Sources/Types/Internal/InternalPart.swift

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,11 @@ struct FileData: Codable, Equatable, Sendable {
3434
self.fileURI = fileURI
3535
self.mimeType = mimeType
3636
}
37+
38+
enum CodingKeys: String, CodingKey {
39+
case fileURI = "fileUri"
40+
case mimeType
41+
}
3742
}
3843

3944
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)

FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ final class IntegrationTests: XCTestCase {
4141

4242
var vertex: VertexAI!
4343
var model: GenerativeModel!
44+
var imagenModel: ImagenModel!
4445
var storage: Storage!
4546
var userID1 = ""
4647

@@ -60,6 +61,9 @@ final class IntegrationTests: XCTestCase {
6061
toolConfig: .init(functionCallingConfig: .none()),
6162
systemInstruction: systemInstruction
6263
)
64+
imagenModel = vertex.imagenModel(
65+
modelName: "imagen-3.0-fast-generate-001"
66+
)
6367

6468
storage = Storage.storage()
6569
}
@@ -235,6 +239,30 @@ final class IntegrationTests: XCTestCase {
235239
XCTAssertTrue(String(describing: error).contains("Firebase App Check token is invalid"))
236240
}
237241
}
242+
243+
// MARK: - Imagen
244+
245+
func testGenerateImage_inlineData() async throws {
246+
let imagePrompt = """
247+
A realistic photo of a male lion, mane thick and dark, standing proudly on a rocky outcrop
248+
overlooking a vast African savanna at sunset. Golden hour light, long shadows, sharp focus on
249+
the lion, shallow depth of field, detailed fur texture, DSLR, 85mm lens.
250+
"""
251+
252+
let imageResponse = try await imagenModel.generateImages(prompt: imagePrompt)
253+
254+
XCTAssertNil(imageResponse.raiFilteredReason)
255+
XCTAssertEqual(imageResponse.images.count, 1)
256+
let image = try XCTUnwrap(imageResponse.images.first)
257+
258+
let textResponse = try await model.generateContent(
259+
InlineDataPart(data: image.data, mimeType: "image/png"),
260+
"What is the name of this animal? Answer with the animal name only."
261+
)
262+
263+
let text = try XCTUnwrap(textResponse.text).trimmingCharacters(in: .whitespacesAndNewlines)
264+
XCTAssertEqual(text, "Lion")
265+
}
238266
}
239267

240268
extension StorageReference {

0 commit comments

Comments
 (0)