From 975151cfff6830c5b95ce53aed3b39ff3841acd2 Mon Sep 17 00:00:00 2001 From: Andrew Heard Date: Thu, 27 Jun 2024 11:47:32 -0400 Subject: [PATCH] Add code snippets for text generation --- Package.swift | 5 ++ samples/APIKey.swift | 62 +++++++++++++++++++++++++ samples/TextGeneration.swift | 90 ++++++++++++++++++++++++++++++++++++ 3 files changed, 157 insertions(+) create mode 100644 samples/APIKey.swift create mode 100644 samples/TextGeneration.swift diff --git a/Package.swift b/Package.swift index 8402e52..2a4cb43 100644 --- a/Package.swift +++ b/Package.swift @@ -44,5 +44,10 @@ let package = Package( .process("GoogleAITests/GenerateContentResponses"), ] ), + .testTarget( + name: "CodeSnippetTests", + dependencies: ["GoogleGenerativeAI"], + path: "samples" + ), ] ) diff --git a/samples/APIKey.swift b/samples/APIKey.swift new file mode 100644 index 0000000..3368712 --- /dev/null +++ b/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 + } +} diff --git a/samples/TextGeneration.swift b/samples/TextGeneration.swift new file mode 100644 index 0000000..adf97b4 --- /dev/null +++ b/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] + } +}