From 0adb052a7a32dea8fd8db2b1338f3c39cb517c60 Mon Sep 17 00:00:00 2001
From: Morgan Chen <morganchen12@gmail.com>
Date: Tue, 20 Feb 2024 15:38:37 -0800
Subject: [PATCH] code feedback changes

---
 Sources/GoogleAI/Chat.swift            |  8 +++----
 Sources/GoogleAI/GenerativeModel.swift | 33 ++++++++++++++------------
 2 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/Sources/GoogleAI/Chat.swift b/Sources/GoogleAI/Chat.swift
index d0c6fff..c7cfb85 100644
--- a/Sources/GoogleAI/Chat.swift
+++ b/Sources/GoogleAI/Chat.swift
@@ -88,11 +88,11 @@ public class Chat {
   /// - Parameter content: The new content to send as a single chat message.
   /// - Returns: A stream containing the model's response or an error if an error occurred.
   @available(macOS 12.0, *)
-  public func sendMessageStream(_ contentClosure: @autoclosure () throws -> [ModelContent])
+  public func sendMessageStream(_ content: @autoclosure () throws -> [ModelContent])
     -> AsyncThrowingStream<GenerateContentResponse, Error> {
-    let content: [ModelContent]
+    let resolvedContent: [ModelContent]
     do {
-      content = try contentClosure()
+      resolvedContent = try content()
     } catch let underlying {
       return AsyncThrowingStream { continuation in
         let error: Error
@@ -110,7 +110,7 @@ public class Chat {
         var aggregatedContent: [ModelContent] = []
 
         // Ensure that the new content has the role set.
-        let newContent: [ModelContent] = content.map(populateContentRole(_:))
+        let newContent: [ModelContent] = resolvedContent.map(populateContentRole(_:))
 
         // Send the history alongside the new message as context.
         let request = history + newContent
diff --git a/Sources/GoogleAI/GenerativeModel.swift b/Sources/GoogleAI/GenerativeModel.swift
index 2a95112..2a963f7 100644
--- a/Sources/GoogleAI/GenerativeModel.swift
+++ b/Sources/GoogleAI/GenerativeModel.swift
@@ -111,18 +111,21 @@ public final class GenerativeModel {
   /// - Parameter content: The input(s) given to the model as a prompt.
   /// - Returns: The generated content response from the model.
   /// - Throws: A ``GenerateContentError`` if the request failed.
-  public func generateContent(_ content: [ModelContent]) async throws
+  public func generateContent(_ content: @autoclosure () throws -> [ModelContent]) async throws
     -> GenerateContentResponse {
-    let generateContentRequest = GenerateContentRequest(model: modelResourceName,
-                                                        contents: content,
-                                                        generationConfig: generationConfig,
-                                                        safetySettings: safetySettings,
-                                                        isStreaming: false,
-                                                        options: requestOptions)
     let response: GenerateContentResponse
     do {
+      let generateContentRequest = try GenerateContentRequest(model: modelResourceName,
+                                                              contents: content(),
+                                                              generationConfig: generationConfig,
+                                                              safetySettings: safetySettings,
+                                                              isStreaming: false,
+                                                              options: requestOptions)
       response = try await generativeAIService.loadRequest(request: generateContentRequest)
     } catch {
+      if let imageError = error as? ImageConversionError {
+        throw GenerateContentError.promptImageContentError(underlying: imageError)
+      }
       throw GenerativeModel.generateContentError(from: error)
     }
 
@@ -251,16 +254,16 @@ public final class GenerativeModel {
   /// - Parameter content: The input given to the model as a prompt.
   /// - Returns: The results of running the model's tokenizer on the input; contains
   /// ``CountTokensResponse/totalTokens``.
-  /// - Throws: A ``CountTokensError`` if the tokenization request failed.
-  public func countTokens(_ content: [ModelContent]) async throws
+  /// - Throws: A ``CountTokensError`` if the tokenization request failed or the input content was
+  /// invalid.
+  public func countTokens(_ content: @autoclosure () throws -> [ModelContent]) async throws
     -> CountTokensResponse {
-    let countTokensRequest = CountTokensRequest(
-      model: modelResourceName,
-      contents: content,
-      options: requestOptions
-    )
-
     do {
+      let countTokensRequest = try CountTokensRequest(
+        model: modelResourceName,
+        contents: content(),
+        options: requestOptions
+      )
       return try await generativeAIService.loadRequest(request: countTokensRequest)
     } catch {
       throw CountTokensError.internalError(underlying: error)