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
google-gemini · Aug 16, 2024 · a8edb40 · a8edb40
1 parent db311dd
commit a8edb40
Show file tree

Hide file tree

Showing 3 changed files with 172 additions and 5 deletions.
diff --git a/.github/workflows/samples.yaml b/.github/workflows/samples.yaml
@@ -68,12 +68,15 @@ jobs:
             echo "Testing $file"
             if [[ -f ${file} ]]; then
               # File exists, so needs to be listed.
+              echo $(basename $file)
+              name=$(basename $file)
               if ! grep -q $name ${README}; then
                 echo "Error: Sample not listed in README ($name)"
                 exit 1
               fi
             else
               # File does not exist, ensure it's not listed
+              name=$(basename $file)
               if grep -q $name ${README}; then
                 echo "Error: Sample should not be listed in README ($name)"
                 exit 1

diff --git a/samples/rest/count_tokens.sh b/samples/rest/count_tokens.sh
@@ -4,6 +4,7 @@ SCRIPT_DIR=$(dirname "$0")
 MEDIA_DIR=$(realpath ${SCRIPT_DIR}/../../third_party)
 
 TEXT_PATH=${MEDIA_DIR}/poem.txt
+A11_PATH=${MEDIA_DIR}/a11.txt
 IMG_PATH=${MEDIA_DIR}/organ.jpg
 AUDIO_PATH=${MEDIA_DIR}/sample.mp3
 VIDEO_PATH=${MEDIA_DIR}/Big_Buck_Bunny.mp4
@@ -16,6 +17,13 @@ else
   B64FLAGS="-w0"
 fi
 
+echo "[START tokens_context_window]"
+# [START tokens_context_window]
+curl https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro?key=$GOOGLE_API_KEY > model.json
+jq .inputTokenLimit model.json
+jq .outputTokenLimit model.json
+# [END tokens_context_window]
+
 echo "[START tokens_text_only]"
 # [START tokens_text_only]
 curl https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:countTokens?key=$GOOGLE_API_KEY \
@@ -97,7 +105,6 @@ curl "${upload_url}" \
   --data-binary "@${IMG_PATH}" 2> /dev/null > file_info.json
 
 file_uri=$(jq ".file.uri" file_info.json)
-echo file_uri=$file_uri
 
 curl "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:countTokens?key=$GOOGLE_API_KEY" \
     -H 'Content-Type: application/json' \
@@ -143,13 +150,10 @@ curl "${upload_url}" \
   --data-binary "@${VIDEO_PATH}" 2> /dev/null > file_info.json
 
 file_uri=$(jq ".file.uri" file_info.json)
-echo file_uri=$file_uri
 
 state=$(jq ".file.state" file_info.json)
-echo state=$state
 
 name=$(jq ".file.name" file_info.json)
-echo name=$name
 
 while [[ "($state)" = *"PROCESSING"* ]];
 do
@@ -170,4 +174,115 @@ curl "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:c
           {"file_data":{"mime_type": "video/mp4", "file_uri": '$file_uri'}}]
         }]
        }'
-# [END tokens_multimodal_video_audio_file_api]
+# [END tokens_multimodal_video_audio_file_api]
+
+echo "[START tokens_cached_content]"
+# [START tokens_cached_content]
+echo '{
+  "model": "models/gemini-1.5-flash-001",
+  "contents":[
+    {
+      "parts":[
+        {
+          "inline_data": {
+            "mime_type":"text/plain",
+            "data": "'$(base64 $B64FLAGS $A11_PATH)'"
+          }
+        }
+      ],
+    "role": "user"
+    }
+  ],
+  "systemInstruction": {
+    "parts": [
+      {
+        "text": "You are an expert at analyzing transcripts."
+      }
+    ]
+  },
+  "ttl": "300s"
+}' > request.json
+
+curl -X POST "https://generativelanguage.googleapis.com/v1beta/cachedContents?key=$GOOGLE_API_KEY" \
+ -H 'Content-Type: application/json' \
+ -d @request.json \
+ > cache.json
+
+jq .usageMetadata.totalTokenCount cache.json
+# [END tokens_cached_content]
+
+echo "[START tokens_system_instruction]"
+# [START tokens_system_instruction]
+curl "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro-latest:generateContent?key=$GOOGLE_API_KEY" \
+-H 'Content-Type: application/json' \
+-d '{ "system_instruction": {
+    "parts":
+      { "text": "You are a cat. Your name is Neko."}},
+    "contents": {
+      "parts": {
+        "text": "Hello there"}}}' > system_instructions.json
+
+jq .usageMetadata.totalTokenCount system_instructions.json
+# [END tokens_system_instruction]
+
+echo "[START tokens_tools]"
+# [START tokens_tools]
+cat > tools.json << EOF
+{
+  "function_declarations": [
+    {
+      "name": "enable_lights",
+      "description": "Turn on the lighting system.",
+      "parameters": { "type": "object" }
+    },
+    {
+      "name": "set_light_color",
+      "description": "Set the light color. Lights must be enabled for this to work.",
+      "parameters": {
+        "type": "object",
+        "properties": {
+          "rgb_hex": {
+            "type": "string",
+            "description": "The light color as a 6-digit hex string, e.g. ff0000 for red."
+          }
+        },
+        "required": [
+          "rgb_hex"
+        ]
+      }
+    },
+    {
+      "name": "stop_lights",
+      "description": "Turn off the lighting system.",
+      "parameters": { "type": "object" }
+    }
+  ]
+} 
+EOF
+
+curl "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro-latest:generateContent?key=$GOOGLE_API_KEY" \
+  -H 'Content-Type: application/json' \
+  -d '
+  {
+    "system_instruction": {
+      "parts": {
+        "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."
+      }
+    },
+    "tools": ['$(source "$tools")'],
+
+    "tool_config": {
+      "function_calling_config": {"mode": "none"}
+    },
+
+    "contents": {
+      "role": "user",
+      "parts": {
+        "text": "What can you do?"
+      }
+    }
+  }
+' > tools_output.json
+
+jq .usageMetadata.totalTokenCount tools_output.json
+# [END tokens_tools]
diff --git a/samples/rest/files.sh b/samples/rest/files.sh
@@ -8,6 +8,7 @@ IMG_PATH=${MEDIA_DIR}/organ.jpg
 IMG_PATH_2=${MEDIA_DIR}/Cajun_instruments.jpg
 AUDIO_PATH=${MEDIA_DIR}/sample.mp3
 VIDEO_PATH=${MEDIA_DIR}/Big_Buck_Bunny.mp4
+PDF_PATH=${MEDIA_DIR}/test.pdf
 
 BASE_URL="https://generativelanguage.googleapis.com"
 
@@ -243,6 +244,54 @@ echo
 jq ".candidates[].content.parts[].text" response.json
 # [END files_create_video]
 
+echo "[START files_create_pdf]"
+# [START files_create_pdf]
+NUM_BYTES=$(wc -c < "${PDF_PATH}")
+DISPLAY_NAME=TEXT
+tmp_header_file=upload-header.tmp
+
+# Initial resumable request defining metadata.
+# The upload url is in the response headers dump them to a file.
+curl "${BASE_URL}/upload/v1beta/files?key=${GOOGLE_API_KEY}" \
+  -D upload-header.tmp \
+  -H "X-Goog-Upload-Protocol: resumable" \
+  -H "X-Goog-Upload-Command: start" \
+  -H "X-Goog-Upload-Header-Content-Length: ${NUM_BYTES}" \
+  -H "X-Goog-Upload-Header-Content-Type: application/pdf" \
+  -H "Content-Type: application/json" \
+  -d "{'file': {'display_name': '${DISPLAY_NAME}'}}" 2> /dev/null
+
+upload_url=$(grep -i "x-goog-upload-url: " "${tmp_header_file}" | cut -d" " -f2 | tr -d "\r")
+rm "${tmp_header_file}"
+
+# Upload the actual bytes.
+curl "${upload_url}" \
+  -H "Content-Length: ${NUM_BYTES}" \
+  -H "X-Goog-Upload-Offset: 0" \
+  -H "X-Goog-Upload-Command: upload, finalize" \
+  --data-binary "@${PDF_PATH}" 2> /dev/null > file_info.json
+
+file_uri=$(jq ".file.uri" file_info.json)
+echo file_uri=$file_uri
+
+# Now generate content using that file
+curl "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key=$GOOGLE_API_KEY" \
+    -H 'Content-Type: application/json' \
+    -X POST \
+    -d '{
+      "contents": [{
+        "parts":[
+          {"text": "Can you add a few more lines to this poem?"},
+          {"file_data":{"mime_type": "application/pdf", "file_uri": '$file_uri'}}]
+        }]
+       }' 2> /dev/null > response.json
+
+cat response.json
+echo
+
+jq ".candidates[].content.parts[].text" response.json
+# [END files_create_pdf]
+
 echo "[START files_list]"
 # [START files_list]
 echo "My files: "