Skip to content
Merged
20 changes: 20 additions & 0 deletions FirebaseVertexAI/Sources/GenerateContentResponse.swift
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,26 @@ public struct GenerateContentResponse: Sendable {
}
}

/// Returns inline data parts found in any `Part`s of the first candidate of the response, if any.
public var inlineDataParts: [InlineDataPart] {
guard let candidate = candidates.first else {
VertexLog.error(
code: .generateContentResponseNoCandidates,
"Could not get inline data parts from a response that had no candidates."
)
return []
}
let inlineData: [InlineDataPart] = candidate.content.parts.compactMap { part in
switch part {
case let inlineDataPart as InlineDataPart:
return inlineDataPart
default:
return nil
}
}
return inlineData
}

/// Initializer for SwiftUI previews or tests.
public init(candidates: [Candidate], promptFeedback: PromptFeedback? = nil,
usageMetadata: UsageMetadata? = nil) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,10 @@ struct GenerateContentIntegrationTests {
let candidate = try #require(response.candidates.first)
let inlineDataPart = try #require(candidate.content.parts
.first { $0 is InlineDataPart } as? InlineDataPart)
let inlineDataPartsViaAccessor = response.inlineDataParts
#expect(inlineDataPartsViaAccessor.count == 1)
let inlineDataPartViaAccessor = try #require(inlineDataPartsViaAccessor.first)
#expect(inlineDataPart == inlineDataPartViaAccessor)
#expect(inlineDataPart.mimeType == "image/png")
#expect(inlineDataPart.data.count > 0)
#if canImport(UIKit)
Expand Down
91 changes: 91 additions & 0 deletions FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -1538,6 +1538,97 @@ final class GenerativeModelTests: XCTestCase {
XCTAssertEqual(response.totalTokens, 6)
}

// MARK: - GenerateContentResponse Computed Properties

func testGenerateContentResponse_inlineDataParts_success() throws {
// 1. Create mock parts
let imageData = Data("sample image data".utf8) // Placeholder data
let inlineDataPart = InlineDataPart(data: imageData, mimeType: "image/png")
let textPart = TextPart("This is the text part.")

// 2. Create ModelContent
let modelContent = ModelContent(parts: [textPart, inlineDataPart]) // Mixed parts

// 3. Create Candidate
let candidate = Candidate(
content: modelContent,
safetyRatings: [], // Assuming negligible for this test
finishReason: .stop,
citationMetadata: nil
)

// 4. Create GenerateContentResponse
let response = GenerateContentResponse(candidates: [candidate])

// 5. Assertions for inlineDataParts
let inlineParts = response.inlineDataParts
XCTAssertFalse(inlineParts.isEmpty, "inlineDataParts should not be empty.")
XCTAssertEqual(inlineParts.count, 1, "There should be exactly one InlineDataPart.")

let firstInlinePart = try XCTUnwrap(inlineParts.first, "Could not get the first inline part.")
XCTAssertEqual(firstInlinePart.mimeType, "image/png", "MimeType should match.")
XCTAssertFalse(firstInlinePart.data.isEmpty, "Inline data should not be empty.")
XCTAssertEqual(firstInlinePart.data, imageData) // Verify data content

// 6. Assertion for text (ensure other properties still work)
XCTAssertEqual(response.text, "This is the text part.")

// 7. Assertion for function calls (ensure it's empty)
XCTAssertTrue(response.functionCalls.isEmpty, "functionCalls should be empty.")
}

func testGenerateContentResponse_inlineDataParts_noInlineData() throws {
// 1. Create mock parts (only text)
let textPart = TextPart("This is the text part.")
let funcCallPart = FunctionCallPart(name: "testFunc", args: [:]) // Add another part type

// 2. Create ModelContent
let modelContent = ModelContent(parts: [textPart, funcCallPart])

// 3. Create Candidate
let candidate = Candidate(
content: modelContent,
safetyRatings: [],
finishReason: .stop,
citationMetadata: nil
)

// 4. Create GenerateContentResponse
let response = GenerateContentResponse(candidates: [candidate])

// 5. Assertions for inlineDataParts
let inlineParts = response.inlineDataParts
XCTAssertTrue(inlineParts.isEmpty, "inlineDataParts should be empty.")

// 6. Assertion for text
XCTAssertEqual(response.text, "This is the text part.")

// 7. Assertion for function calls
XCTAssertEqual(response.functionCalls.count, 1)
XCTAssertEqual(response.functionCalls.first?.name, "testFunc")
}

func testGenerateContentResponse_inlineDataParts_noCandidates() throws {
// 1. Create GenerateContentResponse with no candidates
let response = GenerateContentResponse(candidates: [])

// 2. Assertions for inlineDataParts
let inlineParts = response.inlineDataParts
XCTAssertTrue(
inlineParts.isEmpty,
"inlineDataParts should be empty when there are no candidates."
)

// 3. Assertion for text
XCTAssertNil(response.text, "Text should be nil when there are no candidates.")

// 4. Assertion for function calls
XCTAssertTrue(
response.functionCalls.isEmpty,
"functionCalls should be empty when there are no candidates."
)
}

// MARK: - Helpers

private func testFirebaseInfo(appCheck: AppCheckInterop? = nil,
Expand Down
Loading