diff --git a/FirebaseAI/Sources/GenerativeModel.swift b/FirebaseAI/Sources/GenerativeModel.swift index 8d3f5e043a7..e827d8acc3c 100644 --- a/FirebaseAI/Sources/GenerativeModel.swift +++ b/FirebaseAI/Sources/GenerativeModel.swift @@ -137,6 +137,12 @@ public final class GenerativeModel: Sendable { return try await generateContent([ModelContent(parts: parts)]) } + public func generateContent(@ModelContentBuilder _ contentBuilder: () + -> ModelContent) async throws -> GenerateContentResponse { + let content = contentBuilder() + return try await generateContent([content]) + } + /// Generates new content from input content given to the model as a prompt. /// /// - Parameter content: The input(s) given to the model as a prompt. diff --git a/FirebaseAI/Sources/ModelContent.swift b/FirebaseAI/Sources/ModelContent.swift index 7d82bd76445..c9d4a965e75 100644 --- a/FirebaseAI/Sources/ModelContent.swift +++ b/FirebaseAI/Sources/ModelContent.swift @@ -169,3 +169,55 @@ extension ModelContent.InternalPart: Codable { } } } + +// MARK: - ModelContentBuilder + +@resultBuilder +public struct ModelContentBuilder { + typealias Expression = PartsRepresentable + typealias Component = [PartsRepresentable] + typealias Result = ModelContent + + public static func buildExpression(_ expression: PartsRepresentable) -> [any PartsRepresentable] { + return [expression] + } + + public static func buildBlock(_ components: [any PartsRepresentable]...) + -> [any PartsRepresentable] { + return components + } + + public static func buildEither(first component: [any PartsRepresentable]) + -> [any PartsRepresentable] { + return component + } + + public static func buildEither(second component: [any PartsRepresentable]) + -> [any PartsRepresentable] { + return component + } + + public static func buildArray(_ components: [any PartsRepresentable]) + -> [any PartsRepresentable] { + return components + } + + public static func buildArray(_ components: [[any PartsRepresentable]]) + -> [any PartsRepresentable] { + return components.flatMap { $0 } + } + + public static func buildOptional(_ component: [any PartsRepresentable]?) + -> [any PartsRepresentable] { + return component ?? [] + } + + public static func buildLimitedAvailability(_ component: [any PartsRepresentable]) + -> [any PartsRepresentable] { + return component + } + + public static func buildFinalResult(_ component: [any PartsRepresentable]) -> ModelContent { + return ModelContent(parts: component) + } +} diff --git a/FirebaseAI/Tests/Unit/PartsRepresentableTests.swift b/FirebaseAI/Tests/Unit/PartsRepresentableTests.swift index e7531d1da9e..3139c08140f 100644 --- a/FirebaseAI/Tests/Unit/PartsRepresentableTests.swift +++ b/FirebaseAI/Tests/Unit/PartsRepresentableTests.swift @@ -27,6 +27,33 @@ import XCTest @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) final class PartsRepresentableTests: XCTestCase { + func testModelContentBuilderBuildsFromParts() { + let testConditional = false + @ModelContentBuilder func builderTester() -> ModelContent { + "A string prompt of some kind" + + if testConditional { + "A single-conditional string prompt" + } + + if Int.random(in: 0 ..< 5) > 2 { + "A true-branch conditional string prompt" + } else { + "A false-branch conditional string prompt" + } + + for _ in 0 ..< 10 { + "A looped string prompt" + } + + "Finally, a non-conditional string prompt" + } + + let content = builderTester() + XCTAssert(content.parts.count == 13, + "Expected 14 parts, got \(content.parts.count): \(content.parts)") + } + #if !os(watchOS) func testModelContentFromCGImageIsNotEmpty() throws { // adapted from https://forums.swift.org/t/creating-a-cgimage-from-color-array/18634/2