Skip to content

Commit

Permalink
Add an integration test and fix uncovered camelCase JSON issues
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewheard committed Dec 6, 2024
1 parent ee3fde0 commit 996cd4b
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 1 deletion.
15 changes: 14 additions & 1 deletion FirebaseVertexAI/Sources/GenerationConfig.swift
Original file line number Diff line number Diff line change
Expand Up @@ -162,4 +162,17 @@ public struct GenerationConfig {
// MARK: - Codable Conformances

@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
extension GenerationConfig: Encodable {}
extension GenerationConfig: Encodable {
enum CodingKeys: String, CodingKey {
case temperature
case topP
case topK
case candidateCount
case maxOutputTokens
case presencePenalty
case frequencyPenalty
case stopSequences
case responseMIMEType = "responseMimeType"
case responseSchema
}
}
5 changes: 5 additions & 0 deletions FirebaseVertexAI/Sources/Types/Internal/InternalPart.swift
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,11 @@ struct FileData: Codable, Equatable, Sendable {
self.fileURI = fileURI
self.mimeType = mimeType
}

enum CodingKeys: String, CodingKey {
case fileURI = "fileUri"
case mimeType
}
}

@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ final class IntegrationTests: XCTestCase {

var vertex: VertexAI!
var model: GenerativeModel!
var imagenModel: ImagenModel!
var storage: Storage!
var userID1 = ""

Expand All @@ -60,6 +61,9 @@ final class IntegrationTests: XCTestCase {
toolConfig: .init(functionCallingConfig: .none()),
systemInstruction: systemInstruction
)
imagenModel = vertex.imagenModel(
modelName: "imagen-3.0-fast-generate-001"
)

storage = Storage.storage()
}
Expand Down Expand Up @@ -235,6 +239,30 @@ final class IntegrationTests: XCTestCase {
XCTAssertTrue(String(describing: error).contains("Firebase App Check token is invalid"))
}
}

// MARK: - Imagen

func testGenerateImage_inlineData() async throws {
let imagePrompt = """
A realistic photo of a male lion, mane thick and dark, standing proudly on a rocky outcrop
overlooking a vast African savanna at sunset. Golden hour light, long shadows, sharp focus on
the lion, shallow depth of field, detailed fur texture, DSLR, 85mm lens.
"""

let imageResponse = try await imagenModel.generateImages(prompt: imagePrompt)

XCTAssertNil(imageResponse.raiFilteredReason)
XCTAssertEqual(imageResponse.images.count, 1)
let image = try XCTUnwrap(imageResponse.images.first)

let textResponse = try await model.generateContent(
InlineDataPart(data: image.data, mimeType: "image/png"),
"What is the name of this animal? Answer with the animal name only."
)

let text = try XCTUnwrap(textResponse.text).trimmingCharacters(in: .whitespacesAndNewlines)
XCTAssertEqual(text, "Lion")
}
}

extension StorageReference {
Expand Down

0 comments on commit 996cd4b

Please sign in to comment.