Skip to content

Commit 124584b

Browse files
authored
Add count tokens snippets for system instructions and tools (#189)
1 parent ce5715a commit 124584b

File tree

1 file changed

+57
-0
lines changed

1 file changed

+57
-0
lines changed

samples/CountTokens.swift

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,4 +98,61 @@ final class CountTokensSnippets: XCTestCase {
9898
// [END tokens_multimodal_image_inline]
9999
}
100100
#endif // canImport(UIKit)
101+
102+
func testCountTokensSystemInstruction() async throws {
103+
// [START tokens_system_instruction]
104+
let generativeModel =
105+
GenerativeModel(
106+
// Specify a model that supports system instructions, like a Gemini 1.5 model
107+
name: "gemini-1.5-flash",
108+
// Access your API key from your on-demand resource .plist file (see "Set up your API key"
109+
// above)
110+
apiKey: APIKey.default,
111+
systemInstruction: ModelContent(role: "system", parts: "You are a cat. Your name is Neko.")
112+
)
113+
114+
let prompt = "What is your name?"
115+
116+
let response = try await generativeModel.countTokens(prompt)
117+
print("Total Tokens: \(response.totalTokens)")
118+
// [END tokens_system_instruction]
119+
}
120+
121+
func testCountTokensTools() async throws {
122+
// [START tokens_tools]
123+
let generativeModel =
124+
GenerativeModel(
125+
// Specify a model that supports system instructions, like a Gemini 1.5 model
126+
name: "gemini-1.5-flash",
127+
// Access your API key from your on-demand resource .plist file (see "Set up your API key"
128+
// above)
129+
apiKey: APIKey.default,
130+
tools: [Tool(functionDeclarations: [
131+
FunctionDeclaration(
132+
name: "controlLight",
133+
description: "Set the brightness and color temperature of a room light.",
134+
parameters: [
135+
"brightness": Schema(
136+
type: .number,
137+
format: "double",
138+
description: "Light level from 0 to 100. Zero is off and 100 is full brightness."
139+
),
140+
"colorTemperature": Schema(
141+
type: .string,
142+
format: "enum",
143+
description: "Color temperature of the light fixture.",
144+
enumValues: ["daylight", "cool", "warm"]
145+
),
146+
],
147+
requiredParameters: ["brightness", "colorTemperature"]
148+
),
149+
])]
150+
)
151+
152+
let prompt = "Dim the lights so the room feels cozy and warm."
153+
154+
let response = try await generativeModel.countTokens(prompt)
155+
print("Total Tokens: \(response.totalTokens)")
156+
// [END tokens_tools]
157+
}
101158
}

0 commit comments

Comments
 (0)