Add code snippets for text generation

Author: andrewheard

Package.swift

@@ -44,5 +44,10 @@ let package = Package(
         .process("GoogleAITests/GenerateContentResponses"),
       ]
     ),
+    .testTarget(
+      name: "CodeSnippetTests",
+      dependencies: ["GoogleGenerativeAI"],
+      path: "samples"
+    ),
   ]
 )

samples/APIKey.swift

@@ -0,0 +1,62 @@
+// 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
+
+/// A private wrapper for `APIKey`, hiding it from test files.
+private enum APIKeyCodeSnippet {
+  // The implementation of `APIKey` for use in documentation code snippets; shown in
+  // https://ai.google.dev/gemini-api/docs/quickstart?lang=swift
+  // [START setup_api_key]
+  enum APIKey {
+    // Fetch the API key from `GenerativeAI-Info.plist`
+    static var `default`: String {
+      guard let filePath = Bundle.main.path(forResource: "GenerativeAI-Info", ofType: "plist")
+      else {
+        fatalError("Couldn't find file 'GenerativeAI-Info.plist'.")
+      }
+      let plist = NSDictionary(contentsOfFile: filePath)
+      guard let value = plist?.object(forKey: "API_KEY") as? String else {
+        fatalError("Couldn't find key 'API_KEY' in 'GenerativeAI-Info.plist'.")
+      }
+      if value.starts(with: "_") {
+        fatalError(
+          "Follow the instructions at https://ai.google.dev/tutorials/setup to get an API key."
+        )
+      }
+      return value
+    }
+  }
+  // [END setup_api_key]
+}
+
+/// Protocol to ensure that the `APIKey` APIs do not diverge.
+protocol APIKeyProtocol {
+  static var `default`: String { get }
+}
+
+extension APIKeyCodeSnippet.APIKey: APIKeyProtocol {}
+
+/// An implementation of `APIKey` for use in code snippet tests only.
+enum APIKey: APIKeyProtocol {
+  static let apiKeyEnvVar = "API_KEY"
+
+  static var `default`: String {
+    guard let apiKey = ProcessInfo.processInfo.environment[apiKeyEnvVar] else {
+      return ""
+    }
+    return apiKey
+  }
+}

samples/TextGeneration.swift

@@ -0,0 +1,90 @@
+// 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 GoogleGenerativeAI
+import XCTest
+
+@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
+final class TextGeneration: XCTestCase {
+  override func setUpWithError() throws {
+    try XCTSkipIf(APIKey.default.isEmpty, "`\(APIKey.apiKeyEnvVar)` environment variable not set.")
+  }
+
+  func testTextOnlyPrompt() async throws {
+    // [START text_gen_text_only_prompt]
+    let generativeModel =
+      GenerativeModel(
+        // Specify a Gemini model appropriate for your use case
+        name: "gemini-1.5-flash",
+        // Access your API key from your on-demand resource .plist file (see "Set up your API key"
+        // above)
+        apiKey: APIKey.default
+      )
+
+    let prompt = "Write a story about a magic backpack."
+    let response = try await generativeModel.generateContent(prompt)
+    if let text = response.text {
+      print(text)
+    }
+    // [END text_gen_text_only_prompt]
+  }
+
+  func testTextOnlyPromptStreaming() async throws {
+    // [START text_gen_text_only_prompt_streaming]
+    let generativeModel =
+      GenerativeModel(
+        // Specify a Gemini model appropriate for your use case
+        name: "gemini-1.5-flash",
+        // Access your API key from your on-demand resource .plist file (see "Set up your API key"
+        // above)
+        apiKey: APIKey.default
+      )
+
+    let prompt = "Write a story about a magic backpack."
+    // Use streaming with text-only input
+    for try await response in generativeModel.generateContentStream(prompt) {
+      if let text = response.text {
+        print(text)
+      }
+    }
+    // [END text_gen_text_only_prompt_streaming]
+  }
+
+  func testMultimodalOneImagePrompt() async throws {
+    // [START text_gen_multimodal_one_image_prompt]
+    let generativeModel =
+      GenerativeModel(
+        // Specify a Gemini model appropriate for your use case
+        name: "gemini-1.5-flash",
+        // Access your API key from your on-demand resource .plist file (see "Set up your API key"
+        // above)
+        apiKey: APIKey.default
+      )
+
+    #if canImport(UIKit)
+      guard let image = UIImage(systemName: "cloud.sun") else { fatalError() }
+    #elseif canImport(AppKit)
+      guard let image = NSImage(systemSymbolName: "cloud.sun", accessibilityDescription: nil)
+      else { fatalError() }
+    #endif
+
+    let prompt = "What's in this picture?"
+
+    let response = try await generativeModel.generateContent(image, prompt)
+    if let text = response.text {
+      print(text)
+    }
+    // [END text_gen_multimodal_one_image_prompt]
+  }
+}