@@ -98,4 +98,61 @@ final class CountTokensSnippets: XCTestCase {
98
98
// [END tokens_multimodal_image_inline]
99
99
}
100
100
#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
+ }
101
158
}
0 commit comments