Skip to content

Commit

Permalink
Add code snippets for text generation
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewheard committed Jun 27, 2024
1 parent 5478400 commit 975151c
Show file tree
Hide file tree
Showing 3 changed files with 157 additions and 0 deletions.
5 changes: 5 additions & 0 deletions Package.swift
Original file line number Diff line number Diff line change
Expand Up @@ -44,5 +44,10 @@ let package = Package(
.process("GoogleAITests/GenerateContentResponses"),
]
),
.testTarget(
name: "CodeSnippetTests",
dependencies: ["GoogleGenerativeAI"],
path: "samples"
),
]
)
62 changes: 62 additions & 0 deletions samples/APIKey.swift
Original file line number Diff line number Diff line change
@@ -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
}
}
90 changes: 90 additions & 0 deletions samples/TextGeneration.swift
Original file line number Diff line number Diff line change
@@ -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]
}
}

0 comments on commit 975151c

Please sign in to comment.