Skip to content

Commit a8edb40

Browse files
shilpakancharlashilpakancharla
authored
Add additional functions to files & count_tokens. (#490)
* Add PDF file function * Add rest of count_tokens examples * Add context window count * Tested and fixed count_tokens.sh * Updated files.sh to simplify pdf * Update samples.yaml to get basename of file * Update samples.yaml to get basename of file * Update samples.yaml to get basename of file * Update samples.yaml to get basename of file * test samples.yaml
1 parent db311dd commit a8edb40

File tree

3 files changed

+172
-5
lines changed

3 files changed

+172
-5
lines changed

‎.github/workflows/samples.yaml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,12 +68,15 @@ jobs:
6868
echo "Testing $file"
6969
if [[ -f ${file} ]]; then
7070
# File exists, so needs to be listed.
71+
echo $(basename $file)
72+
name=$(basename $file)
7173
if ! grep -q $name ${README}; then
7274
echo "Error: Sample not listed in README ($name)"
7375
exit 1
7476
fi
7577
else
7678
# File does not exist, ensure it's not listed
79+
name=$(basename $file)
7780
if grep -q $name ${README}; then
7881
echo "Error: Sample should not be listed in README ($name)"
7982
exit 1

‎samples/rest/count_tokens.sh

Lines changed: 120 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ SCRIPT_DIR=$(dirname "$0")
44
MEDIA_DIR=$(realpath ${SCRIPT_DIR}/../../third_party)
55

66
TEXT_PATH=${MEDIA_DIR}/poem.txt
7+
A11_PATH=${MEDIA_DIR}/a11.txt
78
IMG_PATH=${MEDIA_DIR}/organ.jpg
89
AUDIO_PATH=${MEDIA_DIR}/sample.mp3
910
VIDEO_PATH=${MEDIA_DIR}/Big_Buck_Bunny.mp4
@@ -16,6 +17,13 @@ else
1617
B64FLAGS="-w0"
1718
fi
1819

20+
echo "[START tokens_context_window]"
21+
# [START tokens_context_window]
22+
curl https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro?key=$GOOGLE_API_KEY > model.json
23+
jq .inputTokenLimit model.json
24+
jq .outputTokenLimit model.json
25+
# [END tokens_context_window]
26+
1927
echo "[START tokens_text_only]"
2028
# [START tokens_text_only]
2129
curl https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:countTokens?key=$GOOGLE_API_KEY \
@@ -97,7 +105,6 @@ curl "${upload_url}" \
97105
--data-binary "@${IMG_PATH}" 2> /dev/null > file_info.json
98106

99107
file_uri=$(jq ".file.uri" file_info.json)
100-
echo file_uri=$file_uri
101108

102109
curl "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:countTokens?key=$GOOGLE_API_KEY" \
103110
-H 'Content-Type: application/json' \
@@ -143,13 +150,10 @@ curl "${upload_url}" \
143150
--data-binary "@${VIDEO_PATH}" 2> /dev/null > file_info.json
144151

145152
file_uri=$(jq ".file.uri" file_info.json)
146-
echo file_uri=$file_uri
147153

148154
state=$(jq ".file.state" file_info.json)
149-
echo state=$state
150155

151156
name=$(jq ".file.name" file_info.json)
152-
echo name=$name
153157

154158
while [[ "($state)" = *"PROCESSING"* ]];
155159
do
@@ -170,4 +174,115 @@ curl "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:c
170174
{"file_data":{"mime_type": "video/mp4", "file_uri": '$file_uri'}}]
171175
}]
172176
}'
173-
# [END tokens_multimodal_video_audio_file_api]
177+
# [END tokens_multimodal_video_audio_file_api]
178+
179+
echo "[START tokens_cached_content]"
180+
# [START tokens_cached_content]
181+
echo '{
182+
"model": "models/gemini-1.5-flash-001",
183+
"contents":[
184+
{
185+
"parts":[
186+
{
187+
"inline_data": {
188+
"mime_type":"text/plain",
189+
"data": "'$(base64 $B64FLAGS $A11_PATH)'"
190+
}
191+
}
192+
],
193+
"role": "user"
194+
}
195+
],
196+
"systemInstruction": {
197+
"parts": [
198+
{
199+
"text": "You are an expert at analyzing transcripts."
200+
}
201+
]
202+
},
203+
"ttl": "300s"
204+
}' > request.json
205+
206+
curl -X POST "https://generativelanguage.googleapis.com/v1beta/cachedContents?key=$GOOGLE_API_KEY" \
207+
-H 'Content-Type: application/json' \
208+
-d @request.json \
209+
> cache.json
210+
211+
jq .usageMetadata.totalTokenCount cache.json
212+
# [END tokens_cached_content]
213+
214+
echo "[START tokens_system_instruction]"
215+
# [START tokens_system_instruction]
216+
curl "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro-latest:generateContent?key=$GOOGLE_API_KEY" \
217+
-H 'Content-Type: application/json' \
218+
-d '{ "system_instruction": {
219+
"parts":
220+
{ "text": "You are a cat. Your name is Neko."}},
221+
"contents": {
222+
"parts": {
223+
"text": "Hello there"}}}' > system_instructions.json
224+
225+
jq .usageMetadata.totalTokenCount system_instructions.json
226+
# [END tokens_system_instruction]
227+
228+
echo "[START tokens_tools]"
229+
# [START tokens_tools]
230+
cat > tools.json << EOF
231+
{
232+
"function_declarations": [
233+
{
234+
"name": "enable_lights",
235+
"description": "Turn on the lighting system.",
236+
"parameters": { "type": "object" }
237+
},
238+
{
239+
"name": "set_light_color",
240+
"description": "Set the light color. Lights must be enabled for this to work.",
241+
"parameters": {
242+
"type": "object",
243+
"properties": {
244+
"rgb_hex": {
245+
"type": "string",
246+
"description": "The light color as a 6-digit hex string, e.g. ff0000 for red."
247+
}
248+
},
249+
"required": [
250+
"rgb_hex"
251+
]
252+
}
253+
},
254+
{
255+
"name": "stop_lights",
256+
"description": "Turn off the lighting system.",
257+
"parameters": { "type": "object" }
258+
}
259+
]
260+
}
261+
EOF
262+
263+
curl "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro-latest:generateContent?key=$GOOGLE_API_KEY" \
264+
-H 'Content-Type: application/json' \
265+
-d '
266+
{
267+
"system_instruction": {
268+
"parts": {
269+
"text": "You are a helpful lighting system bot. You can turn lights on and off, and you can set the color. Do not perform any other tasks."
270+
}
271+
},
272+
"tools": ['$(source "$tools")'],
273+
274+
"tool_config": {
275+
"function_calling_config": {"mode": "none"}
276+
},
277+
278+
"contents": {
279+
"role": "user",
280+
"parts": {
281+
"text": "What can you do?"
282+
}
283+
}
284+
}
285+
' > tools_output.json
286+
287+
jq .usageMetadata.totalTokenCount tools_output.json
288+
# [END tokens_tools]

‎samples/rest/files.sh

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ IMG_PATH=${MEDIA_DIR}/organ.jpg
88
IMG_PATH_2=${MEDIA_DIR}/Cajun_instruments.jpg
99
AUDIO_PATH=${MEDIA_DIR}/sample.mp3
1010
VIDEO_PATH=${MEDIA_DIR}/Big_Buck_Bunny.mp4
11+
PDF_PATH=${MEDIA_DIR}/test.pdf
1112

1213
BASE_URL="https://generativelanguage.googleapis.com"
1314

@@ -243,6 +244,54 @@ echo
243244
jq ".candidates[].content.parts[].text" response.json
244245
# [END files_create_video]
245246

247+
echo "[START files_create_pdf]"
248+
# [START files_create_pdf]
249+
NUM_BYTES=$(wc -c < "${PDF_PATH}")
250+
DISPLAY_NAME=TEXT
251+
tmp_header_file=upload-header.tmp
252+
253+
# Initial resumable request defining metadata.
254+
# The upload url is in the response headers dump them to a file.
255+
curl "${BASE_URL}/upload/v1beta/files?key=${GOOGLE_API_KEY}" \
256+
-D upload-header.tmp \
257+
-H "X-Goog-Upload-Protocol: resumable" \
258+
-H "X-Goog-Upload-Command: start" \
259+
-H "X-Goog-Upload-Header-Content-Length: ${NUM_BYTES}" \
260+
-H "X-Goog-Upload-Header-Content-Type: application/pdf" \
261+
-H "Content-Type: application/json" \
262+
-d "{'file': {'display_name': '${DISPLAY_NAME}'}}" 2> /dev/null
263+
264+
upload_url=$(grep -i "x-goog-upload-url: " "${tmp_header_file}" | cut -d" " -f2 | tr -d "\r")
265+
rm "${tmp_header_file}"
266+
267+
# Upload the actual bytes.
268+
curl "${upload_url}" \
269+
-H "Content-Length: ${NUM_BYTES}" \
270+
-H "X-Goog-Upload-Offset: 0" \
271+
-H "X-Goog-Upload-Command: upload, finalize" \
272+
--data-binary "@${PDF_PATH}" 2> /dev/null > file_info.json
273+
274+
file_uri=$(jq ".file.uri" file_info.json)
275+
echo file_uri=$file_uri
276+
277+
# Now generate content using that file
278+
curl "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key=$GOOGLE_API_KEY" \
279+
-H 'Content-Type: application/json' \
280+
-X POST \
281+
-d '{
282+
"contents": [{
283+
"parts":[
284+
{"text": "Can you add a few more lines to this poem?"},
285+
{"file_data":{"mime_type": "application/pdf", "file_uri": '$file_uri'}}]
286+
}]
287+
}' 2> /dev/null > response.json
288+
289+
cat response.json
290+
echo
291+
292+
jq ".candidates[].content.parts[].text" response.json
293+
# [END files_create_pdf]
294+
246295
echo "[START files_list]"
247296
# [START files_list]
248297
echo "My files: "

0 commit comments

Comments
 (0)