Encode model in GenerateContentRequest only when needed

Author: andrewheard

Diff for: Sources/GoogleAI/GenerateContentRequest.swift

@@ -16,8 +16,11 @@ import Foundation
 
 @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
 struct GenerateContentRequest {
-  /// Model name.
+  // Model name.
   let model: String
+  // If true, the `model` field above is encoded in requests; currently only required when nested in
+  // a `CountTokensRequest`.
+  let isModelEncoded: Bool
   let contents: [ModelContent]
   let generationConfig: GenerationConfig?
   let safetySettings: [SafetySetting]?
@@ -39,6 +42,30 @@ extension GenerateContentRequest: Encodable {
     case toolConfig
     case systemInstruction
   }
+
+  func encode(to encoder: any Encoder) throws {
+    var container = encoder.container(keyedBy: CodingKeys.self)
+
+    if isModelEncoded {
+      try container.encode(model, forKey: .model)
+    }
+    try container.encode(contents, forKey: .contents)
+    if let generationConfig {
+      try container.encode(generationConfig, forKey: .generationConfig)
+    }
+    if let safetySettings {
+      try container.encode(safetySettings, forKey: .safetySettings)
+    }
+    if let tools {
+      try container.encode(tools, forKey: .tools)
+    }
+    if let toolConfig {
+      try container.encode(toolConfig, forKey: .toolConfig)
+    }
+    if let systemInstruction {
+      try container.encode(systemInstruction, forKey: .systemInstruction)
+    }
+  }
 }
 
 @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)

Diff for: Sources/GoogleAI/GenerativeModel.swift

@@ -175,15 +175,18 @@ public final class GenerativeModel {
     -> GenerateContentResponse {
     let response: GenerateContentResponse
     do {
-      let generateContentRequest = try GenerateContentRequest(model: modelResourceName,
-                                                              contents: content(),
-                                                              generationConfig: generationConfig,
-                                                              safetySettings: safetySettings,
-                                                              tools: tools,
-                                                              toolConfig: toolConfig,
-                                                              systemInstruction: systemInstruction,
-                                                              isStreaming: false,
-                                                              options: requestOptions)
+      let generateContentRequest = try GenerateContentRequest(
+        model: modelResourceName,
+        isModelEncoded: false,
+        contents: content(),
+        generationConfig: generationConfig,
+        safetySettings: safetySettings,
+        tools: tools,
+        toolConfig: toolConfig,
+        systemInstruction: systemInstruction,
+        isStreaming: false,
+        options: requestOptions
+      )
       response = try await generativeAIService.loadRequest(request: generateContentRequest)
     } catch {
       if let imageError = error as? ImageConversionError {
@@ -249,15 +252,18 @@ public final class GenerativeModel {
       }
     }
 
-    let generateContentRequest = GenerateContentRequest(model: modelResourceName,
-                                                        contents: evaluatedContent,
-                                                        generationConfig: generationConfig,
-                                                        safetySettings: safetySettings,
-                                                        tools: tools,
-                                                        toolConfig: toolConfig,
-                                                        systemInstruction: systemInstruction,
-                                                        isStreaming: true,
-                                                        options: requestOptions)
+    let generateContentRequest = GenerateContentRequest(
+      model: modelResourceName,
+      isModelEncoded: false,
+      contents: evaluatedContent,
+      generationConfig: generationConfig,
+      safetySettings: safetySettings,
+      tools: tools,
+      toolConfig: toolConfig,
+      systemInstruction: systemInstruction,
+      isStreaming: true,
+      options: requestOptions
+    )
 
     var responseIterator = generativeAIService.loadRequestStream(request: generateContentRequest)
       .makeAsyncIterator()
@@ -326,6 +332,7 @@ public final class GenerativeModel {
     -> CountTokensResponse {
     do {
       let generateContentRequest = try GenerateContentRequest(model: modelResourceName,
+                                                              isModelEncoded: true,
                                                               contents: content(),
                                                               generationConfig: generationConfig,
                                                               safetySettings: safetySettings,

Diff for: Tests/GoogleAITests/GenerateContentRequestTests.swift

@@ -36,6 +36,7 @@ final class GenerateContentRequestTests: XCTestCase {
     let content = [ModelContent(role: role, parts: prompt)]
     let request = GenerateContentRequest(
       model: modelName,
+      isModelEncoded: true,
       contents: content,
       generationConfig: GenerationConfig(temperature: 0.5),
       safetySettings: [SafetySetting(
@@ -108,10 +109,11 @@ final class GenerateContentRequestTests: XCTestCase {
     """)
   }
 
-  func testEncodeRequest_optionalFieldsOmitted() throws {
+  func testEncodeRequest_optionalFieldsOmitted_modelNameEncoded() throws {
     let content = [ModelContent(role: role, parts: prompt)]
     let request = GenerateContentRequest(
       model: modelName,
+      isModelEncoded: true,
       contents: content,
       generationConfig: nil,
       safetySettings: nil,
@@ -141,4 +143,38 @@ final class GenerateContentRequestTests: XCTestCase {
     }
     """)
   }
+
+  func testEncodeRequest_optionalFieldsOmitted_modelNameNotEncoded() throws {
+    let content = [ModelContent(role: role, parts: prompt)]
+    let request = GenerateContentRequest(
+      model: modelName,
+      isModelEncoded: false,
+      contents: content,
+      generationConfig: nil,
+      safetySettings: nil,
+      tools: nil,
+      toolConfig: nil,
+      systemInstruction: nil,
+      isStreaming: false,
+      options: RequestOptions()
+    )
+
+    let jsonData = try encoder.encode(request)
+
+    let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8))
+    XCTAssertEqual(json, """
+    {
+      "contents" : [
+        {
+          "parts" : [
+            {
+              "text" : "\(prompt)"
+            }
+          ],
+          "role" : "\(role)"
+        }
+      ]
+    }
+    """)
+  }
 }