[Firebase AI] Replace Any with TemplateInputRepresentable

andrewheard · andrewheard · commit 432d70b759ec · 2025-10-31T16:30:58.000-04:00
diff --git a/FirebaseAI/Sources/TemplateChatSession.swift b/FirebaseAI/Sources/TemplateChatSession.swift
@@ -50,15 +50,14 @@ public final class TemplateChatSession: Sendable {
   /// - Returns: The content generated by the model.
   /// - Throws: A ``GenerateContentError`` if the request failed.
   public func sendMessage(_ content: [ModelContent],
-                          inputs: [String: Any],
+                          inputs: [String: any TemplateInputRepresentable],
                           options: RequestOptions = RequestOptions()) async throws
     -> GenerateContentResponse {
-    let templateInputs = try inputs.mapValues { try TemplateInput(value: $0) }
     let newContent = content.map(populateContentRole)
     let response = try await model.generateContentWithHistory(
       history: _history.history + newContent,
       template: templateID,
-      inputs: templateInputs,
+      inputs: inputs.mapValues { $0.templateInputRepresentation },
       options: options
     )
     _history.append(contentsOf: newContent)
@@ -80,7 +79,7 @@ public final class TemplateChatSession: Sendable {
   /// - Returns: The content generated by the model.
   /// - Throws: A ``GenerateContentError`` if the request failed.
   public func sendMessage(_ message: any PartsRepresentable,
-                          inputs: [String: Any],
+                          inputs: [String: any TemplateInputRepresentable],
                           options: RequestOptions = RequestOptions()) async throws
     -> GenerateContentResponse {
     return try await sendMessage([ModelContent(parts: message.partsValue)],
@@ -101,15 +100,14 @@ public final class TemplateChatSession: Sendable {
   /// - Returns: An `AsyncThrowingStream` that yields `GenerateContentResponse` objects.
   /// - Throws: A ``GenerateContentError`` if the request failed.
   public func sendMessageStream(_ content: [ModelContent],
-                                inputs: [String: Any],
+                                inputs: [String: any TemplateInputRepresentable],
                                 options: RequestOptions = RequestOptions()) throws
     -> AsyncThrowingStream<GenerateContentResponse, Error> {
-    let templateInputs = try inputs.mapValues { try TemplateInput(value: $0) }
     let newContent = content.map(populateContentRole)
     let stream = try model.generateContentStreamWithHistory(
       history: _history.history + newContent,
       template: templateID,
-      inputs: templateInputs,
+      inputs: inputs.mapValues { $0.templateInputRepresentation },
       options: options
     )
     return AsyncThrowingStream { continuation in
@@ -156,7 +154,7 @@ public final class TemplateChatSession: Sendable {
   /// - Returns: An `AsyncThrowingStream` that yields `GenerateContentResponse` objects.
   /// - Throws: A ``GenerateContentError`` if the request failed.
   public func sendMessageStream(_ message: any PartsRepresentable,
-                                inputs: [String: Any],
+                                inputs: [String: any TemplateInputRepresentable],
                                 options: RequestOptions = RequestOptions()) throws
     -> AsyncThrowingStream<GenerateContentResponse, Error> {
     return try sendMessageStream([ModelContent(parts: message.partsValue)],
diff --git a/FirebaseAI/Sources/TemplateGenerativeModel.swift b/FirebaseAI/Sources/TemplateGenerativeModel.swift
@@ -40,14 +40,13 @@ public final class TemplateGenerativeModel: Sendable {
   /// - Returns: The content generated by the model.
   /// - Throws: A ``GenerateContentError`` if the request failed.
   public func generateContent(templateID: String,
-                              inputs: [String: Any],
+                              inputs: [String: any TemplateInputRepresentable],
                               options: RequestOptions = RequestOptions()) async throws
     -> GenerateContentResponse {
-    let templateInputs = try inputs.mapValues { try TemplateInput(value: $0) }
     return try await generateContentWithHistory(
       history: [],
       template: templateID,
-      inputs: templateInputs,
+      inputs: inputs.mapValues { $0.templateInputRepresentation },
       options: options
     )
   }
@@ -90,13 +89,12 @@ public final class TemplateGenerativeModel: Sendable {
   /// - Returns: An `AsyncThrowingStream` that yields `GenerateContentResponse` objects.
   /// - Throws: A ``GenerateContentError`` if the request failed.
   public func generateContentStream(templateID: String,
-                                    inputs: [String: Any],
+                                    inputs: [String: any TemplateInputRepresentable],
                                     options: RequestOptions = RequestOptions()) throws
     -> AsyncThrowingStream<GenerateContentResponse, Error> {
-    let templateInputs = try inputs.mapValues { try TemplateInput(value: $0) }
     let request = TemplateGenerateContentRequest(
       template: templateID,
-      inputs: templateInputs,
+      inputs: inputs.mapValues { $0.templateInputRepresentation },
       history: [],
       projectID: generativeAIService.firebaseInfo.projectID,
       stream: true,
diff --git a/FirebaseAI/Sources/TemplateImagenModel.swift b/FirebaseAI/Sources/TemplateImagenModel.swift
@@ -39,14 +39,13 @@ public final class TemplateImagenModel: Sendable {
   /// - Returns: The images generated by the model.
   /// - Throws: An error if the request failed.
   public func generateImages(templateID: String,
-                             inputs: [String: Any],
+                             inputs: [String: any TemplateInputRepresentable],
                              options: RequestOptions = RequestOptions()) async throws
     -> ImagenGenerationResponse<ImagenInlineImage> {
-    let templateInputs = try inputs.mapValues { try TemplateInput(value: $0) }
     let projectID = generativeAIService.firebaseInfo.projectID
     let request = TemplateImagenGenerationRequest<ImagenInlineImage>(
       template: templateID,
-      inputs: templateInputs,
+      inputs: inputs.mapValues { $0.templateInputRepresentation },
       projectID: projectID,
       apiConfig: apiConfig,
       options: options
diff --git a/FirebaseAI/Sources/TemplateInput.swift b/FirebaseAI/Sources/TemplateInput.swift
@@ -14,53 +14,45 @@
 
 import Foundation
 
-enum TemplateInput: Encodable, Sendable {
-  case string(String)
-  case int(Int)
-  case double(Double)
-  case bool(Bool)
-  case array([TemplateInput])
-  case dictionary([String: TemplateInput])
+public struct TemplateInput: Sendable {
+  let kind: Kind
 
-  init(value: Any) throws {
-    switch value {
-    case let value as String:
-      self = .string(value)
-    case let value as Int:
-      self = .int(value)
-    case let value as Double:
-      self = .double(value)
-    case let value as Float:
-      self = .double(Double(value))
-    case let value as Bool:
-      self = .bool(value)
-    case let value as [Any]:
-      self = try .array(value.map { try TemplateInput(value: $0) })
-    case let value as [String: Any]:
-      self = try .dictionary(value.mapValues { try TemplateInput(value: $0) })
-    default:
-      throw EncodingError.invalidValue(
-        value,
-        EncodingError.Context(codingPath: [], debugDescription: "Invalid value")
-      )
-    }
+  public init(_ input: some TemplateInputRepresentable) {
+    self = .init(kind: input.templateInputRepresentation.kind)
+  }
+
+  init(kind: Kind) {
+    self.kind = kind
   }
 
-  func encode(to encoder: Encoder) throws {
-    var container = encoder.singleValueContainer()
-    switch self {
-    case let .string(value):
-      try container.encode(value)
-    case let .int(value):
-      try container.encode(value)
-    case let .double(value):
-      try container.encode(value)
-    case let .bool(value):
-      try container.encode(value)
-    case let .array(value):
-      try container.encode(value)
-    case let .dictionary(value):
-      try container.encode(value)
+  enum Kind: Encodable, Sendable {
+    case string(String)
+    case int(Int)
+    case double(Double)
+    case bool(Bool)
+    case array([Kind])
+    case dictionary([String: Kind])
+
+    func encode(to encoder: Encoder) throws {
+      var container = encoder.singleValueContainer()
+      switch self {
+      case let .string(value):
+        try container.encode(value)
+      case let .int(value):
+        try container.encode(value)
+      case let .double(value):
+        try container.encode(value)
+      case let .bool(value):
+        try container.encode(value)
+      case let .array(value):
+        try container.encode(value)
+      case let .dictionary(value):
+        try container.encode(value)
+      }
     }
   }
 }
+
+extension TemplateInput: TemplateInputRepresentable {
+  public var templateInputRepresentation: TemplateInput { self }
+}
diff --git a/FirebaseAI/Sources/Types/Public/TemplateInputRepresentable.swift b/FirebaseAI/Sources/Types/Public/TemplateInputRepresentable.swift
@@ -0,0 +1,55 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import Foundation
+
+/// A type that can be represented as a ``TemplateInput``.
+public protocol TemplateInputRepresentable: Encodable, Sendable {
+  var templateInputRepresentation: TemplateInput { get }
+}
+
+extension String: TemplateInputRepresentable {
+  public var templateInputRepresentation: TemplateInput { TemplateInput(kind: .string(self)) }
+}
+
+extension Int: TemplateInputRepresentable {
+  public var templateInputRepresentation: TemplateInput { TemplateInput(kind: .int(self)) }
+}
+
+extension Double: TemplateInputRepresentable {
+  public var templateInputRepresentation: TemplateInput { TemplateInput(kind: .double(self)) }
+}
+
+extension Float: TemplateInputRepresentable {
+  public var templateInputRepresentation: TemplateInput {
+    TemplateInput(kind: .double(Double(self)))
+  }
+}
+
+extension Bool: TemplateInputRepresentable {
+  public var templateInputRepresentation: TemplateInput { TemplateInput(kind: .bool(self)) }
+}
+
+extension Array: TemplateInputRepresentable where Element: TemplateInputRepresentable {
+  public var templateInputRepresentation: TemplateInput {
+    TemplateInput(kind: .array(map { TemplateInput($0).kind }))
+  }
+}
+
+extension Dictionary: TemplateInputRepresentable
+  where Key == String, Value: TemplateInputRepresentable {
+  public var templateInputRepresentation: TemplateInput {
+    TemplateInput(kind: .dictionary(mapValues { TemplateInput($0).kind }))
+  }
+}
