count_tokens/
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///
12/// To run:
13/// ```
14/// API_KEY=[YOUR_API_KEY] RUST_LOG=info cargo run --package google-generative-ai-rs  --example count_tokens
15/// ``
16#[tokio::main]
17async fn main() -> Result<(), Box<dyn std::error::Error>> {
18    env_logger::init();
19
20    let client = Client::new_from_response_type(
21        ResponseType::CountTokens,
22        env::var("API_KEY").unwrap().to_string(),
23    );
24
25    let txt_request = Request {
26        contents: vec![Content {
27            role: Role::User,
28            parts: vec![Part {
29                text: Some("Write a story about a magic backpack.".to_string()),
30                inline_data: None,
31                file_data: None,
32                video_metadata: None,
33            }],
34        }],
35        tools: vec![],
36        safety_settings: vec![],
37        generation_config: None,
38
39        #[cfg(feature = "beta")]
40        system_instruction: None,
41    };
42
43    let response = client.post(30, &txt_request).await?;
44
45    info!("{:#?}", response);
46
47    Ok(())
48}
count_tokens/count_tokens.rs

count_tokens/
count_tokens.rs
