Skip to content

Commit 74d2f2b

Browse files
authored
Add a code snippet for function calling (#187)
1 parent bad1c2b commit 74d2f2b

File tree

1 file changed

+109
-0
lines changed

1 file changed

+109
-0
lines changed

samples/FunctionCalling.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 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+
@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
25+
final class FunctionCallingSnippets: XCTestCase {
26+
override func setUpWithError() throws {
27+
try XCTSkipIf(
28+
APIKey.default.isEmpty,
29+
"`\(APIKey.apiKeyEnvVar)` environment variable not set."
30+
)
31+
}
32+
33+
func testFunctionCalling() async throws {
34+
// [BEGIN function_calling]
35+
// Calls a hypothetical API to control a light bulb and returns the values that were set.
36+
func controlLight(brightness: Double, colorTemperature: String) -> JSONObject {
37+
return ["brightness": .number(brightness), "colorTemperature": .string(colorTemperature)]
38+
}
39+
40+
let generativeModel =
41+
GenerativeModel(
42+
// Use a model that supports function calling, like a Gemini 1.5 model
43+
name: "gemini-1.5-flash",
44+
// Access your API key from your on-demand resource .plist file (see "Set up your API key"
45+
// above)
46+
apiKey: APIKey.default,
47+
tools: [Tool(functionDeclarations: [
48+
FunctionDeclaration(
49+
name: "controlLight",
50+
description: "Set the brightness and color temperature of a room light.",
51+
parameters: [
52+
"brightness": Schema(
53+
type: .number,
54+
format: "double",
55+
description: "Light level from 0 to 100. Zero is off and 100 is full brightness."
56+
),
57+
"colorTemperature": Schema(
58+
type: .string,
59+
format: "enum",
60+
description: "Color temperature of the light fixture.",
61+
enumValues: ["daylight", "cool", "warm"]
62+
),
63+
],
64+
requiredParameters: ["brightness", "colorTemperature"]
65+
),
66+
])]
67+
)
68+
69+
let chat = generativeModel.startChat()
70+
71+
let prompt = "Dim the lights so the room feels cozy and warm."
72+
73+
// Send the message to the model.
74+
let response1 = try await chat.sendMessage(prompt)
75+
76+
// Check if the model responded with a function call.
77+
// For simplicity, this sample uses the first function call found.
78+
guard let functionCall = response1.functionCalls.first else {
79+
fatalError("Model did not respond with a function call.")
80+
}
81+
// Print an error if the returned function was not declared
82+
guard functionCall.name == "controlLight" else {
83+
fatalError("Unexpected function called: \(functionCall.name)")
84+
}
85+
// Verify that the names and types of the parameters match the declaration
86+
guard case let .number(brightness) = functionCall.args["brightness"] else {
87+
fatalError("Missing argument: brightness")
88+
}
89+
guard case let .string(colorTemperature) = functionCall.args["colorTemperature"] else {
90+
fatalError("Missing argument: colorTemperature")
91+
}
92+
93+
// Call the executable function named in the FunctionCall with the arguments specified in the
94+
// FunctionCall and let it call the hypothetical API.
95+
let apiResponse = controlLight(brightness: brightness, colorTemperature: colorTemperature)
96+
97+
// Send the API response back to the model so it can generate a text response that can be
98+
// displayed to the user.
99+
let response2 = try await chat.sendMessage([ModelContent(
100+
role: "function",
101+
parts: [.functionResponse(FunctionResponse(name: "controlLight", response: apiResponse))]
102+
)])
103+
104+
if let text = response2.text {
105+
print(text)
106+
}
107+
// [END function_calling]
108+
}
109+
}

0 commit comments

Comments
 (0)