Преглед изворни кода

Add result builder api to Firebase AI

Morgan Chen пре 9 месеци
родитељ
комит
3b5a377543

+ 6 - 0
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.

+ 52 - 0
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)
+  }
+}

+ 27 - 0
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