use_completions_vertex/
use_completions_vertex.rs

1use anyhow::Result;
2use schemars::JsonSchema;
3use serde::Deserialize;
4use serde::Serialize;
5
6use allms::{llm::GoogleModels, Completions};
7
8mod utils;
9use utils::get_vertex_token;
10
11#[derive(Deserialize, Serialize, JsonSchema, Debug, Clone)]
12struct TranslationResponse {
13    pub spanish: String,
14    pub french: String,
15    pub german: String,
16    pub polish: String,
17}
18
19#[tokio::main]
20async fn main() -> Result<()> {
21    env_logger::init();
22
23    // Get Vertex API authentication token
24    let google_token_str = get_vertex_token().await?;
25
26    // Example context and instructions
27    let instructions =
28        "Translate the following English sentence to all the languages in the response type: Rust is best for working with LLMs";
29
30    // Get answer using Google GeminiPro via Vertex AI
31    let model = GoogleModels::Gemini2_5FlashLite;
32
33    // **Pre-requisite**: GeminiPro request through Vertex AI require `GOOGLE_PROJECT_ID` environment variable defined
34    let gemini_completion =
35        Completions::new(model, &google_token_str, None, None).version("google-vertex");
36
37    match gemini_completion
38        .get_answer::<TranslationResponse>(instructions)
39        .await
40    {
41        Ok(response) => println!("Vertex Gemini response: {:#?}", response),
42        Err(e) => eprintln!("Error: {:?}", e),
43    }
44
45    // Get answer using a fine-tuned model
46
47    // Using a fine-tuned model requires addressing the endpoint directly
48    // Replace env variable with the endpoint ID of the fine-tuned model
49    let fine_tuned_endpoint_id: String =
50        std::env::var("GOOGLE_VERTEX_ENDPOINT_ID").expect("GOOGLE_VERTEX_ENDPOINT_ID not set");
51    let model = GoogleModels::endpoint(&fine_tuned_endpoint_id);
52
53    let gemini_completion =
54        Completions::new(model, &google_token_str, None, None).version("google-vertex");
55
56    match gemini_completion
57        .get_answer::<TranslationResponse>(instructions)
58        .await
59    {
60        Ok(response) => println!("Vertex Gemini response: {:#?}", response),
61        Err(e) => eprintln!("Error: {:?}", e),
62    }
63
64    Ok(())
65}
use_completions_vertex/use_completions_vertex.rs

use_completions_vertex/
use_completions_vertex.rs
