Skip to content

Commit 1f1f308

Browse files
andrewheardandrewheard
authored
[Vertex AI] Add code snippets for use in docs (#13653)
1 parent 78a181d commit 1f1f308

File tree

3 files changed

+250
-0
lines changed

3 files changed

+250
-0
lines changed

‎FirebaseVertexAI/Tests/Unit/Snippets/FirebaseAppSnippetsUtil.swift‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
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 FirebaseCore
16+
import Foundation
17+
import XCTest
18+
19+
extension FirebaseApp {
20+
static let projectIDEnvVar = "PROJECT_ID"
21+
static let appIDEnvVar = "APP_ID"
22+
static let apiKeyEnvVar = "API_KEY"
23+
24+
static func configureForSnippets() throws {
25+
let environment = ProcessInfo.processInfo.environment
26+
guard let projectID = environment[projectIDEnvVar] else {
27+
throw XCTSkip("No Firebase Project ID specified in environment variable \(projectIDEnvVar).")
28+
}
29+
guard let appID = environment[appIDEnvVar] else {
30+
throw XCTSkip("No Google App ID specified in environment variable \(appIDEnvVar).")
31+
}
32+
guard let apiKey = environment[apiKeyEnvVar] else {
33+
throw XCTSkip("No API key specified in environment variable \(apiKeyEnvVar).")
34+
}
35+
36+
let options = FirebaseOptions(googleAppID: appID, gcmSenderID: "")
37+
options.projectID = projectID
38+
options.apiKey = apiKey
39+
40+
FirebaseApp.configure(options: options)
41+
guard FirebaseApp.isDefaultAppConfigured() else {
42+
XCTFail("Default Firebase app not configured.")
43+
return
44+
}
45+
}
46+
}

‎FirebaseVertexAI/Tests/Unit/Snippets/FunctionCallingSnippets.swift‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
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 FirebaseCore
16+
import FirebaseVertexAI
17+
import XCTest
18+
19+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
20+
final class FunctionCallingSnippets: XCTestCase {
21+
override func setUpWithError() throws {
22+
try FirebaseApp.configureForSnippets()
23+
}
24+
25+
override func tearDown() async throws {
26+
if let app = FirebaseApp.app() {
27+
await app.delete()
28+
}
29+
}
30+
31+
func testFunctionCalling() async throws {
32+
// This function calls a hypothetical external API that returns
33+
// a collection of weather information for a given location on a given date.
34+
func fetchWeather(city: String, state: String, date: String) -> JSONObject {
35+
// TODO(developer): Write a standard function that would call an external weather API.
36+
37+
// For demo purposes, this hypothetical response is hardcoded here in the expected format.
38+
return [
39+
"temperature": .number(38),
40+
"chancePrecipitation": .string("56%"),
41+
"cloudConditions": .string("partlyCloudy"),
42+
]
43+
}
44+
45+
let fetchWeatherTool = FunctionDeclaration(
46+
name: "fetchWeather",
47+
description: "Get the weather conditions for a specific city on a specific date.",
48+
parameters: [
49+
"location": .object(
50+
properties: [
51+
"city": .string(description: "The city of the location."),
52+
"state": .string(description: "The US state of the location."),
53+
],
54+
description: """
55+
The name of the city and its state for which to get the weather. Only cities in the
56+
USA are supported.
57+
"""
58+
),
59+
"date": .string(
60+
description: """
61+
The date for which to get the weather. Date must be in the format: YYYY-MM-DD.
62+
"""
63+
),
64+
]
65+
)
66+
67+
// Initialize the Vertex AI service and the generative model.
68+
// Use a model that supports function calling, like a Gemini 1.5 model.
69+
let model = VertexAI.vertexAI().generativeModel(
70+
modelName: "gemini-1.5-flash",
71+
// Provide the function declaration to the model.
72+
tools: [.functionDeclarations([fetchWeatherTool])]
73+
)
74+
75+
let chat = model.startChat()
76+
let prompt = "What was the weather in Boston on October 17, 2024?"
77+
78+
// Send the user's question (the prompt) to the model using multi-turn chat.
79+
let response = try await chat.sendMessage(prompt)
80+
81+
var functionResponses = [FunctionResponsePart]()
82+
83+
// When the model responds with one or more function calls, invoke the function(s).
84+
for functionCall in response.functionCalls {
85+
if functionCall.name == "fetchWeather" {
86+
// TODO(developer): Handle invalid arguments.
87+
guard case let .object(location) = functionCall.args["location"] else { fatalError() }
88+
guard case let .string(city) = location["city"] else { fatalError() }
89+
guard case let .string(state) = location["state"] else { fatalError() }
90+
guard case let .string(date) = functionCall.args["date"] else { fatalError() }
91+
92+
functionResponses.append(FunctionResponsePart(
93+
name: functionCall.name,
94+
response: fetchWeather(city: city, state: state, date: date)
95+
))
96+
}
97+
// TODO(developer): Handle other potential function calls, if any.
98+
}
99+
100+
// Send the response(s) from the function back to the model so that the model can use it
101+
// to generate its final response.
102+
let finalResponse = try await chat.sendMessage(
103+
[ModelContent(role: "function", parts: functionResponses)]
104+
)
105+
106+
// Log the text response.
107+
print(finalResponse.text ?? "No text in response.")
108+
}
109+
}

‎FirebaseVertexAI/Tests/Unit/Snippets/StructuredOutputSnippets.swift‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
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 FirebaseCore
16+
import FirebaseVertexAI
17+
import XCTest
18+
19+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
20+
final class StructuredOutputSnippets: XCTestCase {
21+
override func setUpWithError() throws {
22+
try FirebaseApp.configureForSnippets()
23+
}
24+
25+
override func tearDown() async throws {
26+
if let app = FirebaseApp.app() {
27+
await app.delete()
28+
}
29+
}
30+
31+
func testStructuredOutputJSONBasic() async throws {
32+
// Provide a JSON schema object using a standard format.
33+
// Later, pass this schema object into `responseSchema` in the generation config.
34+
let jsonSchema = Schema.object(
35+
properties: [
36+
"characters": Schema.array(
37+
items: .object(
38+
properties: [
39+
"name": .string(),
40+
"age": .integer(),
41+
"species": .string(),
42+
"accessory": .enumeration(values: ["hat", "belt", "shoes"]),
43+
],
44+
optionalProperties: ["accessory"]
45+
)
46+
),
47+
]
48+
)
49+
50+
// Initialize the Vertex AI service and the generative model.
51+
// Use a model that supports `responseSchema`, like one of the Gemini 1.5 models.
52+
let model = VertexAI.vertexAI().generativeModel(
53+
modelName: "gemini-1.5-flash",
54+
// In the generation config, set the `responseMimeType` to `application/json`
55+
// and pass the JSON schema object into `responseSchema`.
56+
generationConfig: GenerationConfig(
57+
responseMIMEType: "application/json",
58+
responseSchema: jsonSchema
59+
)
60+
)
61+
62+
let prompt = "For use in a children's card game, generate 10 animal-based characters."
63+
64+
let response = try await model.generateContent(prompt)
65+
print(response.text ?? "No text in response.")
66+
}
67+
68+
func testStructuredOutputEnumBasic() async throws {
69+
// Provide an enum schema object using a standard format.
70+
// Later, pass this schema object into `responseSchema` in the generation config.
71+
let enumSchema = Schema.enumeration(values: ["drama", "comedy", "documentary"])
72+
73+
// Initialize the Vertex AI service and the generative model.
74+
// Use a model that supports `responseSchema`, like one of the Gemini 1.5 models.
75+
let model = VertexAI.vertexAI().generativeModel(
76+
modelName: "gemini-1.5-flash",
77+
// In the generation config, set the `responseMimeType` to `text/x.enum`
78+
// and pass the enum schema object into `responseSchema`.
79+
generationConfig: GenerationConfig(
80+
responseMIMEType: "text/x.enum",
81+
responseSchema: enumSchema
82+
)
83+
)
84+
85+
let prompt = """
86+
The film aims to educate and inform viewers about real-life subjects, events, or people.
87+
It offers a factual record of a particular topic by combining interviews, historical footage,
88+
and narration. The primary purpose of a film is to present information and provide insights
89+
into various aspects of reality.
90+
"""
91+
92+
let response = try await model.generateContent(prompt)
93+
print(response.text ?? "No text in response.")
94+
}
95+
}

0 commit comments

Comments
 (0)