Skip to content

Add inlineDataParts accessor for GenerateContentResponse #14755

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Apr 23, 2025
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