rig_vertexai/
completion.rs

1//! All supported models: <https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/gemini>
2
3use super::Client;
4use crate::types::{
5    completion_request::VertexCompletionRequest, completion_response::VertexGenerateContentOutput,
6};
7use rig_core::completion::{
8    CompletionError, CompletionModel as CompletionModelTrait, CompletionRequest,
9    CompletionResponse, GetTokenUsage,
10};
11use rig_core::streaming::StreamingCompletionResponse;
12use serde::{Deserialize, Serialize};
13
14/// `gemini-1.5-pro`
15pub const GEMINI_1_5_PRO: &str = "gemini-1.5-pro";
16/// `gemini-1.5-flash`
17pub const GEMINI_1_5_FLASH: &str = "gemini-1.5-flash";
18/// `gemini-1.5-pro-latest`
19pub const GEMINI_1_5_PRO_LATEST: &str = "gemini-1.5-pro-latest";
20/// `gemini-1.5-flash-latest`
21pub const GEMINI_1_5_FLASH_LATEST: &str = "gemini-1.5-flash-latest";
22/// `gemini-2.0-flash-exp`
23pub const GEMINI_2_0_FLASH_EXP: &str = "gemini-2.0-flash-exp";
24/// `gemini-2.5-flash-lite`
25pub const GEMINI_2_5_FLASH_LITE: &str = "gemini-2.5-flash-lite";
26/// `gemini-2.5-flash`
27pub const GEMINI_2_5_FLASH: &str = "gemini-2.5-flash";
28/// `gemini-2.5-pro`
29pub const GEMINI_2_5_PRO: &str = "gemini-2.5-pro";
30
31#[derive(Clone)]
32pub struct CompletionModel {
33    pub(crate) client: crate::client::Client,
34    pub model: String,
35}
36
37#[derive(Clone, Serialize, Deserialize)]
38pub struct PlaceholderStreamingResponse;
39
40impl GetTokenUsage for PlaceholderStreamingResponse {
41    fn token_usage(&self) -> rig_core::completion::Usage {
42        rig_core::completion::Usage::new()
43    }
44}
45
46impl CompletionModel {
47    pub fn new(client: Client, model: impl Into<String>) -> Self {
48        Self {
49            client,
50            model: model.into(),
51        }
52    }
53
54    pub fn with_model(client: Client, model: &str) -> Self {
55        Self {
56            client,
57            model: model.into(),
58        }
59    }
60
61    fn model_path(&self) -> Result<String, CompletionError> {
62        let project = self.client.project();
63        let location = self.client.location();
64        Ok(format!(
65            "projects/{project}/locations/{location}/publishers/google/models/{}",
66            self.model
67        ))
68    }
69}
70
71impl CompletionModelTrait for CompletionModel {
72    type Response = VertexGenerateContentOutput;
73    type StreamingResponse = PlaceholderStreamingResponse;
74
75    type Client = Client;
76
77    fn make(client: &Self::Client, model: impl Into<String>) -> Self {
78        Self::new(client.clone(), model.into())
79    }
80
81    async fn completion(
82        &self,
83        request: CompletionRequest,
84    ) -> Result<CompletionResponse<Self::Response>, CompletionError> {
85        tracing::debug!(
86            target: "rig_core::vertexai",
87            "Vertex AI completion request: {request:?}"
88        );
89
90        let vertex_request = VertexCompletionRequest(request);
91
92        let contents = vertex_request.contents()?;
93        let generation_config = vertex_request.generation_config()?;
94        let system_instruction = vertex_request.system_instruction();
95        let tools = vertex_request.tools();
96        let tool_config = vertex_request.tool_config();
97        let model_path = self.model_path()?;
98
99        let mut request_builder = self
100            .client
101            .get_inner()
102            .await
103            .map_err(|error| CompletionError::ProviderError(error.to_string()))?
104            .generate_content()
105            .set_model(&model_path)
106            .set_contents(contents);
107
108        if let Some(config) = generation_config {
109            request_builder = request_builder.set_generation_config(config);
110        }
111
112        if let Some(system_instruction) = system_instruction {
113            request_builder = request_builder.set_system_instruction(system_instruction);
114        }
115
116        if let Some(tools) = tools {
117            request_builder = request_builder.set_tools([tools]);
118        }
119
120        if let Some(tool_config) = tool_config {
121            request_builder = request_builder.set_tool_config(tool_config);
122        }
123
124        let response = request_builder.send().await.map_err(rpc_error)?;
125
126        tracing::debug!(
127            target: "rig_core::vertexai",
128            "Vertex AI completion response: {response:?}"
129        );
130
131        let vertex_output = VertexGenerateContentOutput(response);
132        let completion_response = vertex_output.try_into()?;
133
134        Ok(completion_response)
135    }
136
137    async fn stream(
138        &self,
139        _request: CompletionRequest,
140    ) -> Result<StreamingCompletionResponse<Self::StreamingResponse>, CompletionError> {
141        Err(CompletionError::ProviderError(
142            "Streaming is not supported for Vertex AI in this integration".to_string(),
143        ))
144    }
145}
146
147/// Map a failed `send()` RPC into a [`CompletionError`] that preserves the
148/// provider's gRPC error text verbatim.
149///
150/// Vertex AI uses a non-HTTP (gRPC/SDK) transport, so there is no
151/// [`http::StatusCode`] to attach; the error body is preserved via
152/// [`CompletionError::from_provider_body`] (`status: None`) rather than a
153/// Rig-prefixed [`CompletionError::ProviderError`] diagnostic. (The
154/// `get_inner()` client-init failure stays a `ProviderError` because it is a
155/// Rig-side setup failure, not a provider response.)
156///
157/// Note: the SDK does not distinguish a server-returned gRPC error from a
158/// transport/connection failure, so a pure connection error is also preserved
159/// here (`status: None`) rather than gated out as a Rig diagnostic the way
160/// Bedrock's typed service errors are.
161fn rpc_error(error: impl std::fmt::Display) -> CompletionError {
162    CompletionError::from_provider_body(error.to_string())
163}
164
165#[cfg(test)]
166#[allow(clippy::expect_used, clippy::unwrap_used, clippy::panic)]
167mod tests {
168    use super::*;
169
170    // The `send()` RPC error type comes from the `google-cloud-aiplatform-v1`
171    // SDK and is not trivially constructible, so `rpc_error` is generic over
172    // `impl Display` and we pin it here with a representative error string of
173    // its parameter type. This guards against a revert to `ProviderError`,
174    // which would surface the body as `None`.
175    #[test]
176    fn rpc_error_preserves_raw_text_without_http_status() {
177        let raw = "status: Unavailable, message: \"the service is currently unavailable\"";
178
179        let err = rpc_error(raw);
180
181        assert_eq!(err.provider_response_body(), Some(raw));
182        assert_eq!(err.provider_response_status(), None);
183    }
184}
rig_vertexai/completion.rs

rig_vertexai/
completion.rs
