gcp_vertex_ai_generative_language/
lib.rs

1//! The `gcp-vertex-ai-generative-ai-language` crate deals with the language
2//! part of the Vertex AI Generative models.
3//!
4//! An async client library for GCP Vertex AI Generative models
5
6pub mod auth;
7
8pub use auth::Authentication;
9use google::ai::generativelanguage::v1beta2::discuss_service_client::DiscussServiceClient;
10use google::ai::generativelanguage::v1beta2::model_service_client::ModelServiceClient;
11use google::ai::generativelanguage::v1beta2::text_service_client::TextServiceClient;
12use tonic::codegen::http::uri::InvalidUri;
13use tonic::transport::{Certificate, Channel, ClientTlsConfig};
14
15/// Errors that can occur when using [LanguageClient].
16#[derive(thiserror::Error, Debug)]
17pub enum Error {
18    /// Transport error
19    #[error("tonic transport error - {0}")]
20    Tonic(#[from] tonic::transport::Error),
21    /// Invalid URI.
22    #[error("{0}")]
23    InvalidUri(#[from] InvalidUri),
24    /// Service error.
25    #[error("Status: {}", .0.message())]
26    Status(#[from] tonic::Status),
27}
28
29const CERTIFICATES: &str = include_str!("../certs/roots.pem");
30
31/// Credentials to use to connect to the services
32#[derive(Clone)]
33pub enum Credentials {
34    /// API Key - see https://cloud.google.com/docs/authentication/api-keys
35    ApiKey(String),
36    /// No authentication
37    None,
38}
39
40/// google protos.
41#[allow(missing_docs)]
42pub mod google {
43
44    /// google.api protos.
45    pub mod api {
46
47        include!(concat!(env!("OUT_DIR"), "/google.api.rs"));
48    }
49
50    /// google.ai protos.
51    pub mod ai {
52
53        /// google.ai.generativelanguage protos.
54        pub mod generativelanguage {
55
56            /// google.ai.generativelanguage.v1beta2 protos.
57            pub mod v1beta2 {
58
59                include!(concat!(
60                    env!("OUT_DIR"),
61                    "/google.ai.generativelanguage.v1beta2.rs"
62                ));
63            }
64        }
65    }
66}
67
68/// Generative Language client.
69#[derive(Clone)]
70pub struct LanguageClient {
71    /// The Discuss service client. In particular, this client is used for
72    /// [`DiscussServiceClient::count_message_tokens`] and
73    /// [`DiscussServiceClient::generate_message`].
74    pub discuss_service: DiscussServiceClient<
75        tonic::service::interceptor::InterceptedService<Channel, Authentication>,
76    >,
77    /// The Model service client. Notably, this client is used for
78    /// [`ModelServiceClient::list_models`] and
79    /// [`ModelServiceClient::get_model`].
80    pub model_service: ModelServiceClient<
81        tonic::service::interceptor::InterceptedService<Channel, Authentication>,
82    >,
83    /// The Text service client. Notably, this client is used for
84    /// [`TextServiceClient::generate_text`],
85    /// and [`TextServiceClient::embed_text`].
86    pub text_service:
87        TextServiceClient<tonic::service::interceptor::InterceptedService<Channel, Authentication>>,
88}
89
90impl LanguageClient {
91    /// Creates a new LanguageClient.
92    ///
93    /// # Example
94    ///
95    /// ```
96    /// # tokio_test::block_on(async {
97    /// use std::env;
98    ///
99    /// use gcp_vertex_ai_generative_language::LanguageClient;
100    ///
101    /// let creds = gcp_vertex_ai_generative_language::Credentials::ApiKey("my-api-key".to_string());
102    ///
103    /// let mut client = LanguageClient::new(creds).await.unwrap();
104    ///
105    /// # });
106    /// ```
107    pub async fn new(credentials: Credentials) -> Result<Self, Error> {
108        let domain_name = "generativelanguage.googleapis.com".to_string();
109
110        let tls_config = ClientTlsConfig::new()
111            .ca_certificate(Certificate::from_pem(CERTIFICATES))
112            .domain_name(&domain_name);
113
114        let endpoint = format!("https://{endpoint}", endpoint = domain_name);
115
116        let channel = Channel::from_shared(endpoint)?
117            .user_agent("github.com/ssoudan/gcp-vertex-ai-generative-ai")?
118            .tls_config(tls_config)?
119            .connect_lazy();
120
121        Self::from_channel(credentials, channel).await
122    }
123
124    /// Creates a new LanguageClient from a Channel.
125    pub async fn from_channel(
126        credentials: Credentials,
127        channel: Channel,
128    ) -> Result<LanguageClient, Error> {
129        let discuss_service = {
130            let auth = Authentication::build(credentials.clone()).await?;
131            DiscussServiceClient::with_interceptor(channel.clone(), auth)
132        };
133
134        let model_service = {
135            let auth = Authentication::build(credentials.clone()).await?;
136            ModelServiceClient::with_interceptor(channel.clone(), auth)
137        };
138
139        let text_service = {
140            let auth = Authentication::build(credentials).await?;
141            TextServiceClient::with_interceptor(channel, auth)
142        };
143
144        Ok(Self {
145            discuss_service,
146            model_service,
147            text_service,
148        })
149    }
150}
151
152#[cfg(test)]
153mod test;
154
155#[cfg(test)]
156mod common {
157    use std::env;
158
159    use crate::{Credentials, LanguageClient};
160
161    pub(crate) async fn test_client() -> LanguageClient {
162        let api_key = env::var("GOOGLE_API_KEY").expect("GOOGLE_API_KEY must be set");
163
164        LanguageClient::new(Credentials::ApiKey(api_key))
165            .await
166            .unwrap()
167    }
168}
gcp_vertex_ai_generative_language/lib.rs

gcp_vertex_ai_generative_language/
lib.rs
