diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml
deleted file mode 100644
index c076e8c..0000000
--- a/.github/ISSUE_TEMPLATE/bug_report.yml
+++ /dev/null
@@ -1,23 +0,0 @@
-name: Bug report
-description: Use this template to report bugs
-labels: ["type:bug", "component:swift sdk"]
-body:
-  - type: markdown
-    attributes:
-      value: >
-        **Note:** If this is a support question (e.g. _How do I do XYZ?_), please visit the [Build with Google AI Forum](https://discuss.ai.google.dev/). This is a great place to interact with developers, and to learn, share, and support each other.
-  - type: textarea
-    id: description
-    attributes:
-      label: >
-        Description of the bug:
-  - type: textarea
-    id: behavior
-    attributes:
-      label: >
-        Actual vs expected behavior:
-  - type: textarea
-    id: info
-    attributes:
-      label: >
-        Any other information you'd like to share?
diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml
deleted file mode 100644
index c32ae77..0000000
--- a/.github/ISSUE_TEMPLATE/feature_request.yml
+++ /dev/null
@@ -1,23 +0,0 @@
-name: Feature request
-description: Use this template to suggest a new feature
-labels: ["type:feature request", "component:swift sdk"]
-body:
-  - type: markdown
-    attributes:
-      value: >
-        **Note:** If this is a support question (e.g. _How do I do XYZ?_), please visit the [Build with Google AI Forum](https://discuss.ai.google.dev/). This is a great place to interact with developers, and to learn, share, and support each other.
-  - type: textarea
-    id: description
-    attributes:
-      label: >
-        Description of the feature request:
-  - type: textarea
-    id: behavior
-    attributes:
-      label: >
-        What problem are you trying to solve with this feature?
-  - type: textarea
-    id: info
-    attributes:
-      label: >
-        Any other information you'd like to share?
diff --git a/Examples/GenerativeAICLI/Sources/GenerateContent.swift b/Examples/GenerativeAICLI/Sources/GenerateContent.swift
index ab71c43..669dfe4 100644
--- a/Examples/GenerativeAICLI/Sources/GenerateContent.swift
+++ b/Examples/GenerativeAICLI/Sources/GenerateContent.swift
@@ -55,23 +55,7 @@ struct GenerateContent: AsyncParsableCommand {
 
   mutating func run() async throws {
     do {
-      let safetySettings = [SafetySetting(harmCategory: .dangerousContent, threshold: .blockNone)]
-      // Let the server pick the default config.
-      let config = GenerationConfig(
-        temperature: 0.2,
-        topP: 0.1,
-        topK: 16,
-        candidateCount: 1,
-        maxOutputTokens: isStreaming ? nil : 256,
-        stopSequences: nil
-      )
-
-      let model = GenerativeModel(
-        name: modelNameOrDefault(),
-        apiKey: apiKey,
-        generationConfig: config,
-        safetySettings: safetySettings
-      )
+      let model = GenerativeModel(name: modelNameOrDefault(), apiKey: apiKey)
 
       var parts = [ModelContent.Part]()
 
@@ -115,12 +99,10 @@ struct GenerateContent: AsyncParsableCommand {
   }
 
   func modelNameOrDefault() -> String {
-    if let modelName = modelName {
+    if let modelName {
       return modelName
-    } else if imageURL != nil {
-      return "gemini-1.0-pro-vision-latest"
     } else {
-      return "gemini-1.0-pro"
+      return "gemini-1.5-flash-latest"
     }
   }
 }
diff --git a/Examples/GenerativeAISample/ChatSample/Screens/ConversationScreen.swift b/Examples/GenerativeAISample/ChatSample/Screens/ConversationScreen.swift
index d9c9b62..d0bcde3 100644
--- a/Examples/GenerativeAISample/ChatSample/Screens/ConversationScreen.swift
+++ b/Examples/GenerativeAISample/ChatSample/Screens/ConversationScreen.swift
@@ -94,6 +94,8 @@ struct ConversationScreen: View {
   }
 
   private func sendOrStop() {
+    focusedField = nil
+
     if viewModel.busy {
       viewModel.stop()
     } else {
diff --git a/Examples/GenerativeAISample/ChatSample/ViewModels/ConversationViewModel.swift b/Examples/GenerativeAISample/ChatSample/ViewModels/ConversationViewModel.swift
index 84a8f7e..9508c3f 100644
--- a/Examples/GenerativeAISample/ChatSample/ViewModels/ConversationViewModel.swift
+++ b/Examples/GenerativeAISample/ChatSample/ViewModels/ConversationViewModel.swift
@@ -36,7 +36,7 @@ class ConversationViewModel: ObservableObject {
   private var chatTask: Task<Void, Never>?
 
   init() {
-    model = GenerativeModel(name: "gemini-1.0-pro", apiKey: APIKey.default)
+    model = GenerativeModel(name: "gemini-1.5-flash-latest", apiKey: APIKey.default)
     chat = model.startChat()
   }
 
diff --git a/Examples/GenerativeAISample/FunctionCallingSample/Screens/FunctionCallingScreen.swift b/Examples/GenerativeAISample/FunctionCallingSample/Screens/FunctionCallingScreen.swift
index 4848ec5..c7f4dd5 100644
--- a/Examples/GenerativeAISample/FunctionCallingSample/Screens/FunctionCallingScreen.swift
+++ b/Examples/GenerativeAISample/FunctionCallingSample/Screens/FunctionCallingScreen.swift
@@ -65,6 +65,9 @@ struct FunctionCallingScreen: View {
             }
           }
         })
+        .onTapGesture {
+          focusedField = nil
+        }
       }
       InputField("Message...", text: $userPrompt) {
         Image(systemName: viewModel.busy ? "stop.circle.fill" : "arrow.up.circle.fill")
diff --git a/Examples/GenerativeAISample/FunctionCallingSample/ViewModels/FunctionCallingViewModel.swift b/Examples/GenerativeAISample/FunctionCallingSample/ViewModels/FunctionCallingViewModel.swift
index cc569d4..7ebb821 100644
--- a/Examples/GenerativeAISample/FunctionCallingSample/ViewModels/FunctionCallingViewModel.swift
+++ b/Examples/GenerativeAISample/FunctionCallingSample/ViewModels/FunctionCallingViewModel.swift
@@ -39,7 +39,7 @@ class FunctionCallingViewModel: ObservableObject {
 
   init() {
     model = GenerativeModel(
-      name: "gemini-1.0-pro",
+      name: "gemini-1.5-flash-latest",
       apiKey: APIKey.default,
       tools: [Tool(functionDeclarations: [
         FunctionDeclaration(
diff --git a/Examples/GenerativeAISample/GenerativeAIMultimodalSample/Screens/PhotoReasoningScreen.swift b/Examples/GenerativeAISample/GenerativeAIMultimodalSample/Screens/PhotoReasoningScreen.swift
index 98f3275..9302147 100644
--- a/Examples/GenerativeAISample/GenerativeAIMultimodalSample/Screens/PhotoReasoningScreen.swift
+++ b/Examples/GenerativeAISample/GenerativeAIMultimodalSample/Screens/PhotoReasoningScreen.swift
@@ -20,9 +20,17 @@ import SwiftUI
 struct PhotoReasoningScreen: View {
   @StateObject var viewModel = PhotoReasoningViewModel()
 
+  enum FocusedField: Hashable {
+    case message
+  }
+
+  @FocusState
+  var focusedField: FocusedField?
+
   var body: some View {
     VStack {
       MultimodalInputField(text: $viewModel.userInput, selection: $viewModel.selectedItems)
+        .focused($focusedField, equals: .message)
         .onSubmit {
           onSendTapped()
         }
@@ -47,11 +55,16 @@ struct PhotoReasoningScreen: View {
       }
     }
     .navigationTitle("Multimodal sample")
+    .onAppear {
+      focusedField = .message
+    }
   }
 
   // MARK: - Actions
 
   private func onSendTapped() {
+    focusedField = nil
+
     Task {
       await viewModel.reason()
     }
diff --git a/Examples/GenerativeAISample/GenerativeAIMultimodalSample/ViewModels/PhotoReasoningViewModel.swift b/Examples/GenerativeAISample/GenerativeAIMultimodalSample/ViewModels/PhotoReasoningViewModel.swift
index 181c586..3e3fd19 100644
--- a/Examples/GenerativeAISample/GenerativeAIMultimodalSample/ViewModels/PhotoReasoningViewModel.swift
+++ b/Examples/GenerativeAISample/GenerativeAIMultimodalSample/ViewModels/PhotoReasoningViewModel.swift
@@ -44,7 +44,7 @@ class PhotoReasoningViewModel: ObservableObject {
   private var model: GenerativeModel?
 
   init() {
-    model = GenerativeModel(name: "gemini-1.0-pro-vision-latest", apiKey: APIKey.default)
+    model = GenerativeModel(name: "gemini-1.5-flash-latest", apiKey: APIKey.default)
   }
 
   func reason() async {
diff --git a/Examples/GenerativeAISample/GenerativeAITextSample/Screens/SummarizeScreen.swift b/Examples/GenerativeAISample/GenerativeAITextSample/Screens/SummarizeScreen.swift
index 8fbb89f..748c1ad 100644
--- a/Examples/GenerativeAISample/GenerativeAITextSample/Screens/SummarizeScreen.swift
+++ b/Examples/GenerativeAISample/GenerativeAITextSample/Screens/SummarizeScreen.swift
@@ -28,19 +28,23 @@ struct SummarizeScreen: View {
 
   var body: some View {
     VStack {
-      Text("Enter some text, then tap on _Go_ to summarize it.")
-      HStack(alignment: .top) {
-        TextField("Enter text summarize", text: $userInput, axis: .vertical)
-          .textFieldStyle(.roundedBorder)
-          .onSubmit {
+      VStack(alignment: .leading) {
+        Text("Enter some text, then tap on _Go_ to summarize it.")
+          .padding(.horizontal, 6)
+        HStack(alignment: .top) {
+          TextField("Enter text summarize", text: $userInput, axis: .vertical)
+            .focused($focusedField, equals: .message)
+            .textFieldStyle(.roundedBorder)
+            .onSubmit {
+              onSummarizeTapped()
+            }
+          Button("Go") {
             onSummarizeTapped()
           }
-        Button("Go") {
-          onSummarizeTapped()
+          .padding(.top, 4)
         }
-        .padding(.top, 4)
       }
-      .padding([.horizontal, .bottom])
+      .padding(.horizontal, 16)
 
       List {
         HStack(alignment: .top) {
@@ -61,6 +65,8 @@ struct SummarizeScreen: View {
   }
 
   private func onSummarizeTapped() {
+    focusedField = nil
+
     Task {
       await viewModel.summarize(inputText: userInput)
     }
diff --git a/Examples/GenerativeAISample/GenerativeAITextSample/ViewModels/SummarizeViewModel.swift b/Examples/GenerativeAISample/GenerativeAITextSample/ViewModels/SummarizeViewModel.swift
index c49d6f1..08aab40 100644
--- a/Examples/GenerativeAISample/GenerativeAITextSample/ViewModels/SummarizeViewModel.swift
+++ b/Examples/GenerativeAISample/GenerativeAITextSample/ViewModels/SummarizeViewModel.swift
@@ -32,7 +32,7 @@ class SummarizeViewModel: ObservableObject {
   private var model: GenerativeModel?
 
   init() {
-    model = GenerativeModel(name: "gemini-1.0-pro", apiKey: APIKey.default)
+    model = GenerativeModel(name: "gemini-1.5-flash-latest", apiKey: APIKey.default)
   }
 
   func summarize(inputText: String) async {
diff --git a/Examples/GenerativeAISample/GenerativeAIUIComponents/Sources/GenerativeAIUIComponents/InputField.swift b/Examples/GenerativeAISample/GenerativeAIUIComponents/Sources/GenerativeAIUIComponents/InputField.swift
index 3f12ea6..67941c3 100644
--- a/Examples/GenerativeAISample/GenerativeAIUIComponents/Sources/GenerativeAIUIComponents/InputField.swift
+++ b/Examples/GenerativeAISample/GenerativeAIUIComponents/Sources/GenerativeAIUIComponents/InputField.swift
@@ -60,10 +60,10 @@ public struct InputField<Label>: View where Label: View {
         }
 
         Button(action: submit, label: label)
-          .padding(.top, 4)
+          .padding(.bottom, 4)
       }
     }
-    .padding(.horizontal)
+    .padding(8)
   }
 }
 
diff --git a/Examples/GenerativeAISample/GenerativeAIUIComponents/Sources/GenerativeAIUIComponents/MultimodalInputField.swift b/Examples/GenerativeAISample/GenerativeAIUIComponents/Sources/GenerativeAIUIComponents/MultimodalInputField.swift
index 9519fb3..128da4a 100644
--- a/Examples/GenerativeAISample/GenerativeAIUIComponents/Sources/GenerativeAIUIComponents/MultimodalInputField.swift
+++ b/Examples/GenerativeAISample/GenerativeAIUIComponents/Sources/GenerativeAIUIComponents/MultimodalInputField.swift
@@ -69,7 +69,7 @@ public struct MultimodalInputField: View {
         Button(action: showChooseAttachmentTypePicker) {
           Image(systemName: "plus")
         }
-        .padding(.top, 4)
+        .padding(.top, 10)
 
         VStack(alignment: .leading) {
           TextField(
@@ -110,7 +110,7 @@ public struct MultimodalInputField: View {
         Button(action: submit) {
           Text("Go")
         }
-        .padding(.top, 4)
+        .padding(.top, 8)
       }
     }
     .padding(.horizontal)
diff --git a/README.md b/README.md
index aa9e0ec..7de87dd 100644
--- a/README.md
+++ b/README.md
@@ -19,7 +19,7 @@ For example, with just a few lines of code, you can access Gemini's multimodal c
 generate text from text-and-image input:
 
 ```swift
-let model = GenerativeModel(name: "gemini-1.5-pro-latest", apiKey: "YOUR_API_KEY")
+let model = GenerativeModel(name: "gemini-1.5-flash-latest", apiKey: "YOUR_API_KEY")
 let cookieImage = UIImage(...)
 let prompt = "Do these look store-bought or homemade?"
 
diff --git a/Sources/GoogleAI/Chat.swift b/Sources/GoogleAI/Chat.swift
index a83947e..6549df4 100644
--- a/Sources/GoogleAI/Chat.swift
+++ b/Sources/GoogleAI/Chat.swift
@@ -30,7 +30,11 @@ public class Chat {
   /// model. This will be provided to the model for each message sent as context for the discussion.
   public var history: [ModelContent]
 
-  /// See ``sendMessage(_:)-3ify5``.
+  /// Sends a message using the existing history of this chat as context. If successful, the message
+  /// and response will be added to the history. If unsuccessful, history will remain unchanged.
+  /// - Parameter parts: The new content to send as a single chat message.
+  /// - Returns: The model's response if no error occurred.
+  /// - Throws: A ``GenerateContentError`` if an error occurred.
   public func sendMessage(_ parts: any ThrowingPartsRepresentable...) async throws
     -> GenerateContentResponse {
     return try await sendMessage([ModelContent(parts: parts)])
@@ -76,7 +80,10 @@ public class Chat {
     return result
   }
 
-  /// See ``sendMessageStream(_:)-4abs3``.
+  /// Sends a message using the existing history of this chat as context. If successful, the message
+  /// and response will be added to the history. If unsuccessful, history will remain unchanged.
+  /// - Parameter parts: 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(_ parts: any ThrowingPartsRepresentable...)
     -> AsyncThrowingStream<GenerateContentResponse, Error> {
diff --git a/Sources/GoogleAI/CountTokensRequest.swift b/Sources/GoogleAI/CountTokensRequest.swift
index f695ffd..6ab5e30 100644
--- a/Sources/GoogleAI/CountTokensRequest.swift
+++ b/Sources/GoogleAI/CountTokensRequest.swift
@@ -17,7 +17,7 @@ import Foundation
 @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
 struct CountTokensRequest {
   let model: String
-  let contents: [ModelContent]
+  let generateContentRequest: GenerateContentRequest
   let options: RequestOptions
 }
 
@@ -45,7 +45,7 @@ public struct CountTokensResponse {
 @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
 extension CountTokensRequest: Encodable {
   enum CodingKeys: CodingKey {
-    case contents
+    case generateContentRequest
   }
 }
 
diff --git a/Sources/GoogleAI/GenerateContentRequest.swift b/Sources/GoogleAI/GenerateContentRequest.swift
index 40c3fe0..620b659 100644
--- a/Sources/GoogleAI/GenerateContentRequest.swift
+++ b/Sources/GoogleAI/GenerateContentRequest.swift
@@ -31,6 +31,7 @@ struct GenerateContentRequest {
 @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
 extension GenerateContentRequest: Encodable {
   enum CodingKeys: String, CodingKey {
+    case model
     case contents
     case generationConfig
     case safetySettings
diff --git a/Sources/GoogleAI/GenerateContentResponse.swift b/Sources/GoogleAI/GenerateContentResponse.swift
index 683df0c..04c41f7 100644
--- a/Sources/GoogleAI/GenerateContentResponse.swift
+++ b/Sources/GoogleAI/GenerateContentResponse.swift
@@ -45,11 +45,17 @@ public struct GenerateContentResponse {
       Logging.default.error("Could not get text from a response that had no candidates.")
       return nil
     }
-    guard let text = candidate.content.parts.first?.text else {
+    let textValues: [String] = candidate.content.parts.compactMap { part in
+      guard case let .text(text) = part else {
+        return nil
+      }
+      return text
+    }
+    guard textValues.count > 0 else {
       Logging.default.error("Could not get a text part from the first candidate.")
       return nil
     }
-    return text
+    return textValues.joined(separator: " ")
   }
 
   /// Returns function calls found in any `Part`s of the first candidate of the response, if any.
diff --git a/Sources/GoogleAI/GenerationConfig.swift b/Sources/GoogleAI/GenerationConfig.swift
index 850d1e9..c4bd37c 100644
--- a/Sources/GoogleAI/GenerationConfig.swift
+++ b/Sources/GoogleAI/GenerationConfig.swift
@@ -70,6 +70,12 @@ public struct GenerationConfig {
   /// - `application/json`: JSON response in the candidates.
   public let responseMIMEType: String?
 
+  /// Output response schema of the generated candidate text.
+  ///
+  /// - Note: This only applies when the specified ``responseMIMEType`` supports a schema; currently
+  ///   this is limited to `application/json`.
+  public let responseSchema: Schema?
+
   /// Creates a new `GenerationConfig` value.
   ///
   /// - Parameters:
@@ -80,9 +86,11 @@ public struct GenerationConfig {
   ///   - maxOutputTokens: See ``maxOutputTokens``.
   ///   - stopSequences: See ``stopSequences``.
   ///   - responseMIMEType: See ``responseMIMEType``.
+  ///   - responseSchema: See ``responseSchema``.
   public init(temperature: Float? = nil, topP: Float? = nil, topK: Int? = nil,
               candidateCount: Int? = nil, maxOutputTokens: Int? = nil,
-              stopSequences: [String]? = nil, responseMIMEType: String? = nil) {
+              stopSequences: [String]? = nil, responseMIMEType: String? = nil,
+              responseSchema: Schema? = nil) {
     // Explicit init because otherwise if we re-arrange the above variables it changes the API
     // surface.
     self.temperature = temperature
@@ -92,6 +100,7 @@ public struct GenerationConfig {
     self.maxOutputTokens = maxOutputTokens
     self.stopSequences = stopSequences
     self.responseMIMEType = responseMIMEType
+    self.responseSchema = responseSchema
   }
 }
 
diff --git a/Sources/GoogleAI/GenerativeAIRequest.swift b/Sources/GoogleAI/GenerativeAIRequest.swift
index 3072dc0..c605f35 100644
--- a/Sources/GoogleAI/GenerativeAIRequest.swift
+++ b/Sources/GoogleAI/GenerativeAIRequest.swift
@@ -26,9 +26,8 @@ protocol GenerativeAIRequest: Encodable {
 /// Configuration parameters for sending requests to the backend.
 @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
 public struct RequestOptions {
-  /// The request’s timeout interval in seconds; if not specified uses the default value for a
-  /// `URLRequest`.
-  let timeout: TimeInterval?
+  /// The request’s timeout interval in seconds.
+  let timeout: TimeInterval
 
   /// The API version to use in requests to the backend.
   let apiVersion: String
@@ -36,10 +35,9 @@ public struct RequestOptions {
   /// Initializes a request options object.
   ///
   /// - Parameters:
-  ///   - timeout The request’s timeout interval in seconds; if not specified uses the default value
-  ///   for a `URLRequest`.
-  ///   - apiVersion The API version to use in requests to the backend; defaults to "v1beta".
-  public init(timeout: TimeInterval? = nil, apiVersion: String = "v1beta") {
+  ///   - timeout: The request’s timeout interval in seconds; defaults to 300 seconds (5 minutes).
+  ///   - apiVersion: The API version to use in requests to the backend; defaults to "v1beta".
+  public init(timeout: TimeInterval = 300.0, apiVersion: String = "v1beta") {
     self.timeout = timeout
     self.apiVersion = apiVersion
   }
diff --git a/Sources/GoogleAI/GenerativeAIService.swift b/Sources/GoogleAI/GenerativeAIService.swift
index 2c1fa87..d1140af 100644
--- a/Sources/GoogleAI/GenerativeAIService.swift
+++ b/Sources/GoogleAI/GenerativeAIService.swift
@@ -44,9 +44,9 @@ struct GenerativeAIService {
 
     // Verify the status code is 200
     guard response.statusCode == 200 else {
-      Logging.default.error("[GoogleGenerativeAI] The server responded with an error: \(response)")
+      Logging.network.error("[GoogleGenerativeAI] The server responded with an error: \(response)")
       if let responseString = String(data: data, encoding: .utf8) {
-        Logging.network.error("[GoogleGenerativeAI] Response payload: \(responseString)")
+        Logging.default.error("[GoogleGenerativeAI] Response payload: \(responseString)")
       }
 
       throw parseError(responseData: data)
@@ -92,14 +92,14 @@ struct GenerativeAIService {
 
         // Verify the status code is 200
         guard response.statusCode == 200 else {
-          Logging.default
+          Logging.network
             .error("[GoogleGenerativeAI] The server responded with an error: \(response)")
           var responseBody = ""
           for try await line in stream.lines {
             responseBody += line + "\n"
           }
 
-          Logging.network.error("[GoogleGenerativeAI] Response payload: \(responseBody)")
+          Logging.default.error("[GoogleGenerativeAI] Response payload: \(responseBody)")
           continuation.finish(throwing: parseError(responseBody: responseBody))
 
           return
@@ -159,10 +159,7 @@ struct GenerativeAIService {
     let encoder = JSONEncoder()
     encoder.keyEncodingStrategy = .convertToSnakeCase
     urlRequest.httpBody = try encoder.encode(request)
-
-    if let timeoutInterval = request.options.timeout {
-      urlRequest.timeoutInterval = timeoutInterval
-    }
+    urlRequest.timeoutInterval = request.options.timeout
 
     return urlRequest
   }
diff --git a/Sources/GoogleAI/GenerativeAISwift.swift b/Sources/GoogleAI/GenerativeAISwift.swift
index e492076..1e0ccdb 100644
--- a/Sources/GoogleAI/GenerativeAISwift.swift
+++ b/Sources/GoogleAI/GenerativeAISwift.swift
@@ -17,10 +17,11 @@ import Foundation
   #warning("Only iOS, macOS, and Catalyst targets are currently fully supported.")
 #endif
 
-/// Constants associated with the GenerativeAISwift SDK
+/// Constants associated with the GenerativeAISwift SDK.
 @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
 public enum GenerativeAISwift {
   /// String value of the SDK version
-  public static let version = "0.5.2"
+  public static let version = "0.5.4"
+  /// The Google AI backend endpoint URL.
   public static let baseURL = "https://generativelanguage.googleapis.com"
 }
diff --git a/Sources/GoogleAI/GenerativeModel.swift b/Sources/GoogleAI/GenerativeModel.swift
index 1d0b986..3e38251 100644
--- a/Sources/GoogleAI/GenerativeModel.swift
+++ b/Sources/GoogleAI/GenerativeModel.swift
@@ -48,14 +48,14 @@ public final class GenerativeModel {
   /// Initializes a new remote model with the given parameters.
   ///
   /// - Parameters:
-  ///   - name: The name of the model to use, e.g., `"gemini-1.5-pro-latest"`; see
+  ///   - name: The name of the model to use, for example `"gemini-1.5-pro-latest"`; see
   ///     [Gemini models](https://ai.google.dev/models/gemini) for a list of supported model names.
   ///   - apiKey: The API key for your project.
   ///   - generationConfig: The content generation parameters your model should use.
   ///   - safetySettings: A value describing what types of harmful content your model should allow.
   ///   - tools: A list of ``Tool`` objects  that the model may use to generate the next response.
   ///   - systemInstruction: Instructions that direct the model to behave a certain way; currently
-  ///     only text content is supported, e.g.,
+  ///     only text content is supported, for example
   ///     `ModelContent(role: "system", parts: "You are a cat. Your name is Neko.")`.
   ///   - toolConfig: Tool configuration for any `Tool` specified in the request.
   ///   - requestOptions Configuration parameters for sending requests to the backend.
@@ -144,7 +144,7 @@ public final class GenerativeModel {
   /// [zero-shot](https://developers.google.com/machine-learning/glossary/generative#zero-shot-prompting)
   /// or "direct" prompts. For
   /// [few-shot](https://developers.google.com/machine-learning/glossary/generative#few-shot-prompting)
-  /// prompts, see ``generateContent(_:)-58rm0``.
+  /// prompts, see `generateContent(_ content: @autoclosure () throws -> [ModelContent])`.
   ///
   /// - Parameter content: The input(s) given to the model as a prompt (see
   /// ``ThrowingPartsRepresentable``
@@ -203,7 +203,7 @@ public final class GenerativeModel {
   /// [zero-shot](https://developers.google.com/machine-learning/glossary/generative#zero-shot-prompting)
   /// or "direct" prompts. For
   /// [few-shot](https://developers.google.com/machine-learning/glossary/generative#few-shot-prompting)
-  /// prompts, see ``generateContent(_:)-58rm0``.
+  /// prompts, see `generateContent(_ content: @autoclosure () throws -> [ModelContent])`.
   ///
   /// - Parameter content: The input(s) given to the model as a prompt (see
   /// ``ThrowingPartsRepresentable``
@@ -292,7 +292,7 @@ public final class GenerativeModel {
   /// [zero-shot](https://developers.google.com/machine-learning/glossary/generative#zero-shot-prompting)
   /// or "direct" prompts. For
   /// [few-shot](https://developers.google.com/machine-learning/glossary/generative#few-shot-prompting)
-  /// input, see ``countTokens(_:)-9spwl``.
+  /// input, see `countTokens(_ content: @autoclosure () throws -> [ModelContent])`.
   ///
   /// - Parameter content: The input(s) given to the model as a prompt (see
   /// ``ThrowingPartsRepresentable``
@@ -315,9 +315,18 @@ public final class GenerativeModel {
   public func countTokens(_ content: @autoclosure () throws -> [ModelContent]) async throws
     -> CountTokensResponse {
     do {
-      let countTokensRequest = try CountTokensRequest(
+      let generateContentRequest = try GenerateContentRequest(model: modelResourceName,
+                                                              contents: content(),
+                                                              generationConfig: generationConfig,
+                                                              safetySettings: safetySettings,
+                                                              tools: tools,
+                                                              toolConfig: toolConfig,
+                                                              systemInstruction: systemInstruction,
+                                                              isStreaming: false,
+                                                              options: requestOptions)
+      let countTokensRequest = CountTokensRequest(
         model: modelResourceName,
-        contents: content(),
+        generateContentRequest: generateContentRequest,
         options: requestOptions
       )
       return try await generativeAIService.loadRequest(request: countTokensRequest)
@@ -350,7 +359,7 @@ public final class GenerativeModel {
   }
 }
 
-/// See ``GenerativeModel/countTokens(_:)-9spwl``.
+/// An error thrown in `GenerativeModel.countTokens(_:)`.
 @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
 public enum CountTokensError: Error {
   case internalError(underlying: Error)
diff --git a/Tests/GoogleAITests/GenerateContentRequestTests.swift b/Tests/GoogleAITests/GenerateContentRequestTests.swift
new file mode 100644
index 0000000..a808799
--- /dev/null
+++ b/Tests/GoogleAITests/GenerateContentRequestTests.swift
@@ -0,0 +1,144 @@
+// Copyright 2024 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
+import XCTest
+
+@testable import GoogleGenerativeAI
+
+@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
+final class GenerateContentRequestTests: XCTestCase {
+  let encoder = JSONEncoder()
+  let role = "test-role"
+  let prompt = "test-prompt"
+  let modelName = "test-model-name"
+
+  override func setUp() {
+    encoder.outputFormatting = .init(
+      arrayLiteral: .prettyPrinted, .sortedKeys, .withoutEscapingSlashes
+    )
+  }
+
+  // MARK: GenerateContentRequest Encoding
+
+  func testEncodeRequest_allFieldsIncluded() throws {
+    let content = [ModelContent(role: role, parts: prompt)]
+    let request = GenerateContentRequest(
+      model: modelName,
+      contents: content,
+      generationConfig: GenerationConfig(temperature: 0.5),
+      safetySettings: [SafetySetting(
+        harmCategory: .dangerousContent,
+        threshold: .blockLowAndAbove
+      )],
+      tools: [Tool(functionDeclarations: [FunctionDeclaration(
+        name: "test-function-name",
+        description: "test-function-description",
+        parameters: nil
+      )])],
+      toolConfig: ToolConfig(functionCallingConfig: FunctionCallingConfig(mode: .auto)),
+      systemInstruction: ModelContent(role: "system", parts: "test-system-instruction"),
+      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)"
+        }
+      ],
+      "generationConfig" : {
+        "temperature" : 0.5
+      },
+      "model" : "\(modelName)",
+      "safetySettings" : [
+        {
+          "category" : "HARM_CATEGORY_DANGEROUS_CONTENT",
+          "threshold" : "BLOCK_LOW_AND_ABOVE"
+        }
+      ],
+      "systemInstruction" : {
+        "parts" : [
+          {
+            "text" : "test-system-instruction"
+          }
+        ],
+        "role" : "system"
+      },
+      "toolConfig" : {
+        "functionCallingConfig" : {
+          "mode" : "AUTO"
+        }
+      },
+      "tools" : [
+        {
+          "functionDeclarations" : [
+            {
+              "description" : "test-function-description",
+              "name" : "test-function-name",
+              "parameters" : {
+                "type" : "OBJECT"
+              }
+            }
+          ]
+        }
+      ]
+    }
+    """)
+  }
+
+  func testEncodeRequest_optionalFieldsOmitted() throws {
+    let content = [ModelContent(role: role, parts: prompt)]
+    let request = GenerateContentRequest(
+      model: modelName,
+      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)"
+        }
+      ],
+      "model" : "\(modelName)"
+    }
+    """)
+  }
+}
diff --git a/Tests/GoogleAITests/GenerateContentResponses/unary-success-function-call-mixed-content.json b/Tests/GoogleAITests/GenerateContentResponses/unary-success-function-call-mixed-content.json
new file mode 100644
index 0000000..6e7ce27
--- /dev/null
+++ b/Tests/GoogleAITests/GenerateContentResponses/unary-success-function-call-mixed-content.json
@@ -0,0 +1,37 @@
+{
+  "candidates": [
+    {
+      "content": {
+        "parts": [
+          {
+            "text": "The sum of [1, 2,"
+          },
+          {
+            "functionCall": {
+              "name": "sum",
+              "args": {
+                "y": 1,
+                "x": 2
+              }
+            }
+          },
+          {
+            "text": "3] is"
+          },
+          {
+            "functionCall": {
+              "name": "sum",
+              "args": {
+                "y": 3,
+                "x": 3
+              }
+            }
+          }
+        ],
+        "role": "model"
+      },
+      "finishReason": "STOP",
+      "index": 0
+    }
+  ]
+}
diff --git a/Tests/GoogleAITests/GenerateContentResponses/unary-success-function-call-parallel-calls.json b/Tests/GoogleAITests/GenerateContentResponses/unary-success-function-call-parallel-calls.json
new file mode 100644
index 0000000..d535f8e
--- /dev/null
+++ b/Tests/GoogleAITests/GenerateContentResponses/unary-success-function-call-parallel-calls.json
@@ -0,0 +1,40 @@
+{
+  "candidates": [
+    {
+      "content": {
+        "parts": [
+          {
+            "functionCall": {
+              "name": "sum",
+              "args": {
+                "y": 1,
+                "x": 2
+              }
+            }
+          },
+          {
+            "functionCall": {
+              "name": "sum",
+              "args": {
+                "y": 3,
+                "x": 4
+              }
+            }
+          },
+          {
+            "functionCall": {
+              "name": "sum",
+              "args": {
+                "y": 5,
+                "x": 6
+              }
+            }
+          }
+        ],
+        "role": "model"
+      },
+      "finishReason": "STOP",
+      "index": 0
+    }
+  ]
+}
diff --git a/Tests/GoogleAITests/GenerationConfigTests.swift b/Tests/GoogleAITests/GenerationConfigTests.swift
index 1ad3eaf..06dada7 100644
--- a/Tests/GoogleAITests/GenerationConfigTests.swift
+++ b/Tests/GoogleAITests/GenerationConfigTests.swift
@@ -48,7 +48,11 @@ final class GenerationConfigTests: XCTestCase {
     let candidateCount = 2
     let maxOutputTokens = 256
     let stopSequences = ["END", "DONE"]
-    let responseMIMEType = "text/plain"
+    let responseMIMEType = "application/json"
+    let schemaType = DataType.object
+    let fieldName = "test-field"
+    let fieldType = DataType.string
+    let responseSchema = Schema(type: schemaType, properties: [fieldName: Schema(type: fieldType)])
     let generationConfig = GenerationConfig(
       temperature: temperature,
       topP: topP,
@@ -56,7 +60,8 @@ final class GenerationConfigTests: XCTestCase {
       candidateCount: candidateCount,
       maxOutputTokens: maxOutputTokens,
       stopSequences: stopSequences,
-      responseMIMEType: responseMIMEType
+      responseMIMEType: responseMIMEType,
+      responseSchema: responseSchema
     )
 
     let jsonData = try encoder.encode(generationConfig)
@@ -67,6 +72,14 @@ final class GenerationConfigTests: XCTestCase {
       "candidateCount" : \(candidateCount),
       "maxOutputTokens" : \(maxOutputTokens),
       "responseMIMEType" : "\(responseMIMEType)",
+      "responseSchema" : {
+        "properties" : {
+          "\(fieldName)" : {
+            "type" : "\(fieldType.rawValue)"
+          }
+        },
+        "type" : "\(schemaType.rawValue)"
+      },
       "stopSequences" : [
         "END",
         "DONE"
@@ -79,7 +92,7 @@ final class GenerationConfigTests: XCTestCase {
   }
 
   func testEncodeGenerationConfig_responseMIMEType() throws {
-    let mimeType = "image/jpeg"
+    let mimeType = "text/plain"
     let generationConfig = GenerationConfig(responseMIMEType: mimeType)
 
     let jsonData = try encoder.encode(generationConfig)
@@ -91,4 +104,27 @@ final class GenerationConfigTests: XCTestCase {
     }
     """)
   }
+
+  func testEncodeGenerationConfig_responseMIMETypeWithSchema() throws {
+    let mimeType = "application/json"
+    let schemaType = DataType.array
+    let arrayItemType = DataType.integer
+    let schema = Schema(type: schemaType, items: Schema(type: arrayItemType))
+    let generationConfig = GenerationConfig(responseMIMEType: mimeType, responseSchema: schema)
+
+    let jsonData = try encoder.encode(generationConfig)
+
+    let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8))
+    XCTAssertEqual(json, """
+    {
+      "responseMIMEType" : "\(mimeType)",
+      "responseSchema" : {
+        "items" : {
+          "type" : "\(arrayItemType.rawValue)"
+        },
+        "type" : "\(schemaType.rawValue)"
+      }
+    }
+    """)
+  }
 }
diff --git a/Tests/GoogleAITests/GenerativeModelTests.swift b/Tests/GoogleAITests/GenerativeModelTests.swift
index 9ed3401..5a20343 100644
--- a/Tests/GoogleAITests/GenerativeModelTests.swift
+++ b/Tests/GoogleAITests/GenerativeModelTests.swift
@@ -254,6 +254,40 @@ final class GenerativeModelTests: XCTestCase {
     XCTAssertEqual(response.functionCalls, [functionCall])
   }
 
+  func testGenerateContent_success_functionCall_parallelCalls() async throws {
+    MockURLProtocol
+      .requestHandler = try httpRequestHandler(
+        forResource: "unary-success-function-call-parallel-calls",
+        withExtension: "json"
+      )
+
+    let response = try await model.generateContent(testPrompt)
+
+    XCTAssertEqual(response.candidates.count, 1)
+    let candidate = try XCTUnwrap(response.candidates.first)
+    XCTAssertEqual(candidate.content.parts.count, 3)
+    let functionCalls = response.functionCalls
+    XCTAssertEqual(functionCalls.count, 3)
+  }
+
+  func testGenerateContent_success_functionCall_mixedContent() async throws {
+    MockURLProtocol
+      .requestHandler = try httpRequestHandler(
+        forResource: "unary-success-function-call-mixed-content",
+        withExtension: "json"
+      )
+
+    let response = try await model.generateContent(testPrompt)
+
+    XCTAssertEqual(response.candidates.count, 1)
+    let candidate = try XCTUnwrap(response.candidates.first)
+    XCTAssertEqual(candidate.content.parts.count, 4)
+    let functionCalls = response.functionCalls
+    XCTAssertEqual(functionCalls.count, 2)
+    let text = try XCTUnwrap(response.text)
+    XCTAssertEqual(text, "The sum of [1, 2, 3] is")
+  }
+
   func testGenerateContent_usageMetadata() async throws {
     MockURLProtocol
       .requestHandler = try httpRequestHandler(
@@ -577,6 +611,20 @@ final class GenerativeModelTests: XCTestCase {
     XCTAssertEqual(response.candidates.count, 1)
   }
 
+  func testGenerateContent_requestOptions_defaultTimeout() async throws {
+    let expectedTimeout = 300.0 // Default in timeout in RequestOptions()
+    MockURLProtocol
+      .requestHandler = try httpRequestHandler(
+        forResource: "unary-success-basic-reply-short",
+        withExtension: "json",
+        timeout: expectedTimeout
+      )
+
+    let response = try await model.generateContent(testPrompt)
+
+    XCTAssertEqual(response.candidates.count, 1)
+  }
+
   // MARK: - Generate Content (Streaming)
 
   func testGenerateContentStream_failureInvalidAPIKey() async throws {
@@ -933,6 +981,25 @@ final class GenerativeModelTests: XCTestCase {
     XCTAssertEqual(responses, 1)
   }
 
+  func testGenerateContentStream_requestOptions_defaultTimeout() async throws {
+    let expectedTimeout = 300.0 // Default in timeout in RequestOptions()
+    MockURLProtocol
+      .requestHandler = try httpRequestHandler(
+        forResource: "streaming-success-basic-reply-short",
+        withExtension: "txt",
+        timeout: expectedTimeout
+      )
+
+    var responses = 0
+    let stream = model.generateContentStream(testPrompt)
+    for try await content in stream {
+      XCTAssertNotNil(content.text)
+      responses += 1
+    }
+
+    XCTAssertEqual(responses, 1)
+  }
+
   // MARK: - Count Tokens
 
   func testCountTokens_succeeds() async throws {
@@ -985,6 +1052,20 @@ final class GenerativeModelTests: XCTestCase {
     XCTAssertEqual(response.totalTokens, 6)
   }
 
+  func testCountTokens_requestOptions_defaultTimeout() async throws {
+    let expectedTimeout = 300.0
+    MockURLProtocol
+      .requestHandler = try httpRequestHandler(
+        forResource: "success-total-tokens",
+        withExtension: "json",
+        timeout: expectedTimeout
+      )
+
+    let response = try await model.countTokens(testPrompt)
+
+    XCTAssertEqual(response.totalTokens, 6)
+  }
+
   // MARK: - Model Resource Name
 
   func testModelResourceName_noPrefix() async throws {
@@ -1033,8 +1114,8 @@ final class GenerativeModelTests: XCTestCase {
   private func httpRequestHandler(forResource name: String,
                                   withExtension ext: String,
                                   statusCode: Int = 200,
-                                  timeout: TimeInterval = URLRequest
-                                    .defaultTimeoutInterval()) throws -> ((URLRequest) throws -> (
+                                  timeout: TimeInterval = RequestOptions()
+                                    .timeout) throws -> ((URLRequest) throws -> (
     URLResponse,
     AsyncLineSequence<URL.AsyncBytes>?
   )) {