vertex_count_tokens/
vertex_count_tokens.rs

1use std::env;
2
3use google_generative_ai_rs::v1::{
4    api::Client,
5    gemini::{request::Request, Content, Part, ResponseType, Role},
6};
7use log::info;
8
9/// Counts the tokens used in a prompt using the public API and an API key for authn
10/// See: `https://ai.google.dev/tutorials/rest_quickstart#count_tokens`
11/// You'll need to install the GCP cli tools and set up your GCP project and region.
12///
13/// The ensure you locally authenticated with GCP using the following commands:
14/// ```
15/// gcloud init
16/// gcloud auth application-default login
17/// ```
18///
19/// To run:
20/// ```
21/// GCP_REGION_NAME=[THE REGION WHERE YOUR ENDPOINT IS HOSTED] GCP_PROJECT_ID=[YOUR GCP PROJECT_ID] RUST_LOG=info cargo run --package google-generative-ai-rs  --example vertex_count_tokens
22/// ``
23#[tokio::main]
24async fn main() -> Result<(), Box<dyn std::error::Error>> {
25    env_logger::init();
26    let region = env::var("GCP_REGION_NAME").unwrap().to_string();
27    let project_id = env::var("GCP_PROJECT_ID").unwrap().to_string();
28
29    let client = Client::new_from_region_project_id_response_type(
30        region.to_string(),
31        project_id.to_string(),
32        ResponseType::CountTokens,
33    );
34
35    let txt_request = Request {
36        contents: vec![Content {
37            role: Role::User,
38            parts: vec![Part {
39                text: Some("Write a story about a magic backpack.".to_string()),
40                inline_data: None,
41                file_data: None,
42                video_metadata: None,
43            }],
44        }],
45        tools: vec![],
46        safety_settings: vec![],
47        generation_config: None,
48
49        #[cfg(feature = "beta")]
50        system_instruction: None,
51    };
52
53    let response = client.post(30, &txt_request).await?;
54
55    info!("{:#?}", response);
56
57    Ok(())
58}
vertex_count_tokens/vertex_count_tokens.rs

vertex_count_tokens/
vertex_count_tokens.rs
