Skip to content

Commit b72e9ba

Browse files
andrewheardandrewheard
authored
Add code snippets for text generation (#181)
1 parent 2c09198 commit b72e9ba

File tree

3 files changed

+233
-0
lines changed

3 files changed

+233
-0
lines changed

‎Package.swift

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,5 +44,10 @@ let package = Package(
4444
.process("GoogleAITests/GenerateContentResponses"),
4545
]
4646
),
47+
.testTarget(
48+
name: "CodeSnippetTests",
49+
dependencies: ["GoogleGenerativeAI"],
50+
path: "samples"
51+
),
4752
]
4853
)

‎samples/APIKey.swift

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
// Copyright 2024 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
import Foundation
16+
import XCTest
17+
18+
/// A private wrapper for `APIKey`, hiding it from test files.
19+
private enum APIKeyCodeSnippet {
20+
// The implementation of `APIKey` for use in documentation code snippets; shown in
21+
// https://ai.google.dev/gemini-api/docs/quickstart?lang=swift
22+
// [START setup_api_key]
23+
enum APIKey {
24+
// Fetch the API key from `GenerativeAI-Info.plist`
25+
static var `default`: String {
26+
guard let filePath = Bundle.main.path(forResource: "GenerativeAI-Info", ofType: "plist")
27+
else {
28+
fatalError("Couldn't find file 'GenerativeAI-Info.plist'.")
29+
}
30+
let plist = NSDictionary(contentsOfFile: filePath)
31+
guard let value = plist?.object(forKey: "API_KEY") as? String else {
32+
fatalError("Couldn't find key 'API_KEY' in 'GenerativeAI-Info.plist'.")
33+
}
34+
if value.starts(with: "_") {
35+
fatalError(
36+
"Follow the instructions at https://ai.google.dev/tutorials/setup to get an API key."
37+
)
38+
}
39+
return value
40+
}
41+
}
42+
// [END setup_api_key]
43+
}
44+
45+
/// Protocol to ensure that the `APIKey` APIs do not diverge.
46+
protocol APIKeyProtocol {
47+
static var `default`: String { get }
48+
}
49+
50+
extension APIKeyCodeSnippet.APIKey: APIKeyProtocol {}
51+
52+
/// An implementation of `APIKey` for use in code snippet tests only.
53+
enum APIKey: APIKeyProtocol {
54+
static let apiKeyEnvVar = "API_KEY"
55+
56+
static var `default`: String {
57+
guard let apiKey = ProcessInfo.processInfo.environment[apiKeyEnvVar] else {
58+
return ""
59+
}
60+
return apiKey
61+
}
62+
}

‎samples/TextGeneration.swift

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
// Copyright 2024 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
import GoogleGenerativeAI
16+
import XCTest
17+
18+
// Set up your API Key
19+
// ====================
20+
// To use the Gemini API, you'll need an API key. To learn more, see the "Set up your API Key"
21+
// section in the Gemini API quickstart:
22+
// https://ai.google.dev/gemini-api/docs/quickstart?lang=swift#set-up-api-key
23+
24+
#if canImport(UIKit)
25+
@available(iOS 15.0, macCatalyst 15.0, *)
26+
final class TextGeneration: XCTestCase {
27+
override func setUpWithError() throws {
28+
try XCTSkipIf(
29+
APIKey.default.isEmpty,
30+
"`\(APIKey.apiKeyEnvVar)` environment variable not set."
31+
)
32+
}
33+
34+
func testTextOnlyPrompt() async throws {
35+
// [START text_gen_text_only_prompt]
36+
let generativeModel =
37+
GenerativeModel(
38+
// Specify a Gemini model appropriate for your use case
39+
name: "gemini-1.5-flash",
40+
// Access your API key from your on-demand resource .plist file (see "Set up your API key"
41+
// above)
42+
apiKey: APIKey.default
43+
)
44+
45+
let prompt = "Write a story about a magic backpack."
46+
let response = try await generativeModel.generateContent(prompt)
47+
if let text = response.text {
48+
print(text)
49+
}
50+
// [END text_gen_text_only_prompt]
51+
}
52+
53+
func testTextOnlyPromptStreaming() async throws {
54+
// [START text_gen_text_only_prompt_streaming]
55+
let generativeModel =
56+
GenerativeModel(
57+
// Specify a Gemini model appropriate for your use case
58+
name: "gemini-1.5-flash",
59+
// Access your API key from your on-demand resource .plist file (see "Set up your API key"
60+
// above)
61+
apiKey: APIKey.default
62+
)
63+
64+
let prompt = "Write a story about a magic backpack."
65+
// Use streaming with text-only input
66+
for try await response in generativeModel.generateContentStream(prompt) {
67+
if let text = response.text {
68+
print(text)
69+
}
70+
}
71+
// [END text_gen_text_only_prompt_streaming]
72+
}
73+
74+
func testMultimodalOneImagePrompt() async throws {
75+
// [START text_gen_multimodal_one_image_prompt]
76+
let generativeModel =
77+
GenerativeModel(
78+
// Specify a Gemini model appropriate for your use case
79+
name: "gemini-1.5-flash",
80+
// Access your API key from your on-demand resource .plist file (see "Set up your API key"
81+
// above)
82+
apiKey: APIKey.default
83+
)
84+
85+
guard let image = UIImage(systemName: "cloud.sun") else { fatalError() }
86+
87+
let prompt = "What's in this picture?"
88+
89+
let response = try await generativeModel.generateContent(image, prompt)
90+
if let text = response.text {
91+
print(text)
92+
}
93+
// [END text_gen_multimodal_one_image_prompt]
94+
}
95+
96+
func testMultimodalOneImagePromptStreaming() async throws {
97+
// [START text_gen_multimodal_one_image_prompt_streaming]
98+
let generativeModel =
99+
GenerativeModel(
100+
// Specify a Gemini model appropriate for your use case
101+
name: "gemini-1.5-flash",
102+
// Access your API key from your on-demand resource .plist file (see "Set up your API key"
103+
// above)
104+
apiKey: APIKey.default
105+
)
106+
107+
guard let image = UIImage(systemName: "cloud.sun") else { fatalError() }
108+
109+
let prompt = "What's in this picture?"
110+
111+
for try await response in generativeModel.generateContentStream(image, prompt) {
112+
if let text = response.text {
113+
print(text)
114+
}
115+
}
116+
// [END text_gen_multimodal_one_image_prompt_streaming]
117+
}
118+
119+
func testMultimodalMultiImagePrompt() async throws {
120+
// [START text_gen_multimodal_multi_image_prompt]
121+
let generativeModel =
122+
GenerativeModel(
123+
// Specify a Gemini model appropriate for your use case
124+
name: "gemini-1.5-flash",
125+
// Access your API key from your on-demand resource .plist file (see "Set up your API key"
126+
// above)
127+
apiKey: APIKey.default
128+
)
129+
130+
guard let image1 = UIImage(systemName: "cloud.sun") else { fatalError() }
131+
guard let image2 = UIImage(systemName: "cloud.heavyrain") else { fatalError() }
132+
133+
let prompt = "What's the difference between these pictures?"
134+
135+
let response = try await generativeModel.generateContent(image1, image2, prompt)
136+
if let text = response.text {
137+
print(text)
138+
}
139+
// [END text_gen_multimodal_multi_image_prompt]
140+
}
141+
142+
func testMultimodalMultiImagePromptStreaming() async throws {
143+
// [START text_gen_multimodal_multi_image_prompt_streaming]
144+
let generativeModel =
145+
GenerativeModel(
146+
// Specify a Gemini model appropriate for your use case
147+
name: "gemini-1.5-flash",
148+
// Access your API key from your on-demand resource .plist file (see "Set up your API key"
149+
// above)
150+
apiKey: APIKey.default
151+
)
152+
153+
guard let image1 = UIImage(systemName: "cloud.sun") else { fatalError() }
154+
guard let image2 = UIImage(systemName: "cloud.heavyrain") else { fatalError() }
155+
156+
let prompt = "What's the difference between these pictures?"
157+
158+
for try await response in generativeModel.generateContentStream(image1, image2, prompt) {
159+
if let text = response.text {
160+
print(text)
161+
}
162+
}
163+
// [END text_gen_multimodal_multi_image_prompt_streaming]
164+
}
165+
}
166+
#endif // canImport(UIKit)

0 commit comments

Comments
 (0)