uni_xervo/provider/
vertexai.rs

1use crate::api::{ModelAliasSpec, ModelTask};
2use crate::error::{Result, RuntimeError};
3use crate::provider::remote_common::{
4    RemoteProviderBase, build_google_generate_payload, check_http_status,
5};
6use crate::traits::{
7    EmbedResult, EmbeddingModel, GenerationOptions, GenerationResult, GeneratorModel,
8    LoadedModelHandle, Message, ModelProvider, ProviderCapabilities, ProviderHealth, TokenUsage,
9};
10use async_trait::async_trait;
11use reqwest::Client;
12use serde_json::json;
13use std::sync::Arc;
14
15fn options_map<'a>(
16    provider_id: &str,
17    options: &'a serde_json::Value,
18) -> Result<Option<&'a serde_json::Map<String, serde_json::Value>>> {
19    match options {
20        serde_json::Value::Null => Ok(None),
21        serde_json::Value::Object(map) => Ok(Some(map)),
22        _ => Err(RuntimeError::Config(format!(
23            "Options for provider '{}' must be a JSON object or null",
24            provider_id
25        ))),
26    }
27}
28
29fn option_string(
30    provider_id: &str,
31    map: Option<&serde_json::Map<String, serde_json::Value>>,
32    key: &str,
33) -> Result<Option<String>> {
34    let Some(map) = map else {
35        return Ok(None);
36    };
37    let Some(value) = map.get(key) else {
38        return Ok(None);
39    };
40    let s = value.as_str().ok_or_else(|| {
41        RuntimeError::Config(format!(
42            "Option '{}' for provider '{}' must be a string",
43            key, provider_id
44        ))
45    })?;
46    Ok(Some(s.to_string()))
47}
48
49fn option_u32(
50    provider_id: &str,
51    map: Option<&serde_json::Map<String, serde_json::Value>>,
52    key: &str,
53) -> Result<Option<u32>> {
54    let Some(map) = map else {
55        return Ok(None);
56    };
57    let Some(value) = map.get(key) else {
58        return Ok(None);
59    };
60    let n = value.as_u64().ok_or_else(|| {
61        RuntimeError::Config(format!(
62            "Option '{}' for provider '{}' must be a positive integer",
63            key, provider_id
64        ))
65    })?;
66    if n == 0 {
67        return Err(RuntimeError::Config(format!(
68            "Option '{}' for provider '{}' must be greater than 0",
69            key, provider_id
70        )));
71    }
72    let n_u32 = u32::try_from(n).map_err(|_| {
73        RuntimeError::Config(format!(
74            "Option '{}' for provider '{}' is out of range for u32",
75            key, provider_id
76        ))
77    })?;
78    Ok(Some(n_u32))
79}
80
81/// Resolved and validated Vertex AI configuration extracted from a
82/// [`ModelAliasSpec`]'s options and environment variables.
83#[derive(Clone)]
84struct VertexAiResolvedOptions {
85    token: String,
86    project_id: String,
87    location: String,
88    publisher: String,
89    embedding_dimensions: Option<u32>,
90}
91
92impl VertexAiResolvedOptions {
93    /// Vertex AI host. The `global` location uses the unprefixed
94    /// `aiplatform.googleapis.com`; all other locations use the
95    /// `{location}-aiplatform.googleapis.com` regional prefix.
96    fn host(&self) -> String {
97        if self.location == "global" {
98            "aiplatform.googleapis.com".to_string()
99        } else {
100            format!("{}-aiplatform.googleapis.com", self.location)
101        }
102    }
103
104    fn from_spec(spec: &ModelAliasSpec) -> Result<Self> {
105        let provider_id = "remote/vertexai";
106        let map = options_map(provider_id, &spec.options)?;
107
108        let token_env = option_string(provider_id, map, "api_token_env")?
109            .unwrap_or_else(|| "VERTEX_AI_TOKEN".to_string());
110        let token = std::env::var(&token_env)
111            .map_err(|_| RuntimeError::Config(format!("{} env var not set", token_env)))?;
112
113        let project_id = if let Some(project_id) = option_string(provider_id, map, "project_id")? {
114            project_id
115        } else {
116            std::env::var("VERTEX_AI_PROJECT").map_err(|_| {
117                RuntimeError::Config(
118                    "project_id option not set and VERTEX_AI_PROJECT env var not set".to_string(),
119                )
120            })?
121        };
122
123        let location =
124            option_string(provider_id, map, "location")?.unwrap_or_else(|| "us-central1".into());
125        let publisher =
126            option_string(provider_id, map, "publisher")?.unwrap_or_else(|| "google".into());
127        let embedding_dimensions = option_u32(provider_id, map, "embedding_dimensions")?;
128
129        Ok(Self {
130            token,
131            project_id,
132            location,
133            publisher,
134            embedding_dimensions,
135        })
136    }
137}
138
139/// Remote provider that calls the [Google Vertex AI](https://cloud.google.com/vertex-ai/docs)
140/// prediction and generation endpoints for embedding and text generation.
141///
142/// Requires the `VERTEX_AI_TOKEN` environment variable (or a custom env var
143/// via `api_token_env`) and either the `project_id` option or the
144/// `VERTEX_AI_PROJECT` env var.
145pub struct RemoteVertexAIProvider {
146    base: RemoteProviderBase,
147}
148
149impl RemoteVertexAIProvider {
150    pub fn new() -> Self {
151        Self::default()
152    }
153
154    #[cfg(test)]
155    fn insert_test_breaker(&self, key: crate::api::ModelRuntimeKey, age: std::time::Duration) {
156        self.base.insert_test_breaker(key, age);
157    }
158
159    #[cfg(test)]
160    fn breaker_count(&self) -> usize {
161        self.base.breaker_count()
162    }
163
164    #[cfg(test)]
165    fn force_cleanup_now_for_test(&self) {
166        self.base.force_cleanup_now_for_test();
167    }
168}
169
170impl Default for RemoteVertexAIProvider {
171    fn default() -> Self {
172        Self {
173            base: RemoteProviderBase::new(),
174        }
175    }
176}
177
178#[async_trait]
179impl ModelProvider for RemoteVertexAIProvider {
180    fn provider_id(&self) -> &'static str {
181        "remote/vertexai"
182    }
183
184    fn capabilities(&self) -> ProviderCapabilities {
185        ProviderCapabilities {
186            supported_tasks: vec![ModelTask::Embed, ModelTask::Generate],
187        }
188    }
189
190    async fn load(&self, spec: &ModelAliasSpec) -> Result<LoadedModelHandle> {
191        let cb = self.base.circuit_breaker_for(spec);
192        let resolved = VertexAiResolvedOptions::from_spec(spec)?;
193
194        match spec.task {
195            ModelTask::Embed => {
196                let model = VertexAiEmbeddingModel {
197                    client: self.base.client.clone(),
198                    cb: cb.clone(),
199                    model_id: spec.model_id.clone(),
200                    options: resolved.clone(),
201                    dimensions: resolved.embedding_dimensions.unwrap_or(768),
202                };
203                let handle: Arc<dyn EmbeddingModel> = Arc::new(model);
204                Ok(Arc::new(handle) as LoadedModelHandle)
205            }
206            ModelTask::Generate => {
207                let model = VertexAiGeneratorModel {
208                    client: self.base.client.clone(),
209                    cb,
210                    model_id: spec.model_id.clone(),
211                    options: resolved,
212                };
213                let handle: Arc<dyn GeneratorModel> = Arc::new(model);
214                Ok(Arc::new(handle) as LoadedModelHandle)
215            }
216            ModelTask::Raw => Err(RuntimeError::CapabilityMismatch(
217                "Vertex AI provider does not support task Raw".to_string(),
218            )),
219            _ => Err(RuntimeError::CapabilityMismatch(format!(
220                "Vertex AI provider does not support task {:?}",
221                spec.task
222            ))),
223        }
224    }
225
226    async fn health(&self) -> ProviderHealth {
227        ProviderHealth::Healthy
228    }
229}
230
231/// Embedding model backed by the Vertex AI prediction API.
232pub struct VertexAiEmbeddingModel {
233    client: Client,
234    cb: crate::reliability::CircuitBreakerWrapper,
235    model_id: String,
236    options: VertexAiResolvedOptions,
237    dimensions: u32,
238}
239
240impl VertexAiEmbeddingModel {
241    fn endpoint_url(&self) -> String {
242        format!(
243            "https://{}/v1/projects/{}/locations/{}/publishers/{}/models/{}:predict",
244            self.options.host(),
245            self.options.project_id,
246            self.options.location,
247            self.options.publisher,
248            self.model_id
249        )
250    }
251}
252
253#[async_trait]
254impl EmbeddingModel for VertexAiEmbeddingModel {
255    async fn embed(&self, texts: &[&str]) -> Result<EmbedResult> {
256        let texts: Vec<String> = texts.iter().map(|s| s.to_string()).collect();
257
258        self.cb
259            .call(move || async move {
260                let instances: Vec<_> = texts.iter().map(|t| json!({ "content": t })).collect();
261                let response = self
262                    .client
263                    .post(self.endpoint_url())
264                    .header("Authorization", format!("Bearer {}", self.options.token))
265                    .json(&json!({ "instances": instances }))
266                    .send()
267                    .await
268                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
269
270                let body: serde_json::Value = check_http_status("Vertex AI", response)?
271                    .json()
272                    .await
273                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
274
275                let predictions = body
276                    .get("predictions")
277                    .and_then(|v| v.as_array())
278                    .ok_or_else(|| {
279                        RuntimeError::ApiError("Invalid response: missing predictions".to_string())
280                    })?;
281
282                let mut vectors = Vec::new();
283                let mut token_total: u64 = 0;
284                let mut saw_token_count = false;
285                for item in predictions {
286                    let values_opt = item
287                        .get("embeddings")
288                        .and_then(|e| e.get("values").and_then(|v| v.as_array()))
289                        .or_else(|| {
290                            item.get("embeddings")
291                                .and_then(|e| e.as_array())
292                                .or_else(|| item.get("values").and_then(|v| v.as_array()))
293                        });
294
295                    let values = values_opt.ok_or_else(|| {
296                        RuntimeError::ApiError(
297                            "Invalid embedding format in Vertex AI response".to_string(),
298                        )
299                    })?;
300
301                    let vec: Vec<f32> = values
302                        .iter()
303                        .filter_map(|v| v.as_f64().map(|f| f as f32))
304                        .collect();
305                    vectors.push(vec);
306
307                    // Each prediction may carry `embeddings.statistics.token_count`.
308                    if let Some(tc) = item
309                        .get("embeddings")
310                        .and_then(|e| e.get("statistics"))
311                        .and_then(|s| s.get("token_count"))
312                        .and_then(|t| t.as_u64())
313                    {
314                        saw_token_count = true;
315                        token_total += tc;
316                    }
317                }
318
319                // Vertex AI predict reports per-instance token counts under
320                // `embeddings.statistics.token_count`; sum them when present.
321                let usage = if saw_token_count {
322                    Some(TokenUsage {
323                        prompt_tokens: token_total as usize,
324                        completion_tokens: 0,
325                        total_tokens: token_total as usize,
326                    })
327                } else {
328                    None
329                };
330
331                Ok(EmbedResult { vectors, usage })
332            })
333            .await
334    }
335
336    fn dimensions(&self) -> u32 {
337        self.dimensions
338    }
339}
340
341impl crate::traits::ModelInfo for VertexAiEmbeddingModel {
342    fn model_id(&self) -> &str {
343        &self.model_id
344    }
345}
346
347/// Text generation model backed by the Vertex AI `generateContent` endpoint.
348pub struct VertexAiGeneratorModel {
349    client: Client,
350    cb: crate::reliability::CircuitBreakerWrapper,
351    model_id: String,
352    options: VertexAiResolvedOptions,
353}
354
355impl VertexAiGeneratorModel {
356    fn endpoint_url(&self) -> String {
357        format!(
358            "https://{}/v1/projects/{}/locations/{}/publishers/{}/models/{}:generateContent",
359            self.options.host(),
360            self.options.project_id,
361            self.options.location,
362            self.options.publisher,
363            self.model_id
364        )
365    }
366}
367
368impl crate::traits::ModelInfo for VertexAiGeneratorModel {
369    fn model_id(&self) -> &str {
370        &self.model_id
371    }
372}
373
374#[async_trait]
375impl GeneratorModel for VertexAiGeneratorModel {
376    async fn generate(
377        &self,
378        messages: &[Message],
379        options: GenerationOptions,
380    ) -> Result<GenerationResult> {
381        let messages: Vec<Message> = messages.to_vec();
382
383        self.cb
384            .call(move || async move {
385                let payload = build_google_generate_payload(&messages, &options);
386                let response = self
387                    .client
388                    .post(self.endpoint_url())
389                    .header("Authorization", format!("Bearer {}", self.options.token))
390                    .json(&payload)
391                    .send()
392                    .await
393                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
394
395                let body: serde_json::Value = check_http_status("Vertex AI", response)?
396                    .json()
397                    .await
398                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
399
400                let candidates = body
401                    .get("candidates")
402                    .and_then(|v| v.as_array())
403                    .ok_or_else(|| RuntimeError::ApiError("No candidates returned".to_string()))?;
404
405                let first_candidate = candidates
406                    .first()
407                    .ok_or_else(|| RuntimeError::ApiError("Empty candidates".to_string()))?;
408
409                let content_parts = first_candidate
410                    .get("content")
411                    .and_then(|c| c.get("parts"))
412                    .and_then(|p| p.as_array())
413                    .ok_or_else(|| RuntimeError::ApiError("Invalid content format".to_string()))?;
414
415                let text = content_parts
416                    .first()
417                    .and_then(|p| p.get("text"))
418                    .and_then(|t| t.as_str())
419                    .unwrap_or("")
420                    .to_string();
421
422                let usage = body.get("usageMetadata").map(|u| TokenUsage {
423                    prompt_tokens: u["promptTokenCount"].as_u64().unwrap_or(0) as usize,
424                    completion_tokens: u["candidatesTokenCount"].as_u64().unwrap_or(0) as usize,
425                    total_tokens: u["totalTokenCount"].as_u64().unwrap_or(0) as usize,
426                });
427
428                Ok(GenerationResult {
429                    text,
430                    usage,
431                    images: vec![],
432                    audio: None,
433                })
434            })
435            .await
436    }
437}
438
439#[cfg(test)]
440mod tests {
441    use super::*;
442    use crate::api::ModelRuntimeKey;
443    use crate::provider::remote_common::RemoteProviderBase;
444    use crate::traits::ModelProvider;
445    use std::time::Duration;
446
447    static ENV_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
448
449    fn spec(
450        alias: &str,
451        task: ModelTask,
452        model_id: &str,
453        options: serde_json::Value,
454    ) -> ModelAliasSpec {
455        ModelAliasSpec {
456            alias: alias.to_string(),
457            task,
458            provider_id: "remote/vertexai".to_string(),
459            model_id: model_id.to_string(),
460            revision: None,
461            warmup: crate::api::WarmupPolicy::Lazy,
462            required: false,
463            timeout: None,
464            load_timeout: None,
465            retry: None,
466            options,
467        }
468    }
469
470    #[test]
471    fn endpoint_url_uses_regional_host_by_default() {
472        let opts = VertexAiResolvedOptions {
473            token: "t".into(),
474            project_id: "p".into(),
475            location: "us-central1".into(),
476            publisher: "google".into(),
477            embedding_dimensions: None,
478        };
479        let gen_model = VertexAiGeneratorModel {
480            client: Client::new(),
481            cb: RemoteProviderBase::new().circuit_breaker_for(&spec(
482                "g",
483                ModelTask::Generate,
484                "gemini-2.0-flash",
485                serde_json::Value::Null,
486            )),
487            model_id: "gemini-2.0-flash".into(),
488            options: opts.clone(),
489        };
490        assert_eq!(
491            gen_model.endpoint_url(),
492            "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:generateContent"
493        );
494
495        let emb = VertexAiEmbeddingModel {
496            client: Client::new(),
497            cb: RemoteProviderBase::new().circuit_breaker_for(&spec(
498                "e",
499                ModelTask::Embed,
500                "text-embedding-005",
501                serde_json::Value::Null,
502            )),
503            model_id: "text-embedding-005".into(),
504            options: opts,
505            dimensions: 768,
506        };
507        assert_eq!(
508            emb.endpoint_url(),
509            "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/publishers/google/models/text-embedding-005:predict"
510        );
511    }
512
513    #[test]
514    fn endpoint_url_uses_global_host_when_location_is_global() {
515        let opts = VertexAiResolvedOptions {
516            token: "t".into(),
517            project_id: "p".into(),
518            location: "global".into(),
519            publisher: "google".into(),
520            embedding_dimensions: None,
521        };
522        let gen_model = VertexAiGeneratorModel {
523            client: Client::new(),
524            cb: RemoteProviderBase::new().circuit_breaker_for(&spec(
525                "g",
526                ModelTask::Generate,
527                "gemini-2.0-flash",
528                serde_json::Value::Null,
529            )),
530            model_id: "gemini-2.0-flash".into(),
531            options: opts.clone(),
532        };
533        assert_eq!(
534            gen_model.endpoint_url(),
535            "https://aiplatform.googleapis.com/v1/projects/p/locations/global/publishers/google/models/gemini-2.0-flash:generateContent"
536        );
537
538        let emb = VertexAiEmbeddingModel {
539            client: Client::new(),
540            cb: RemoteProviderBase::new().circuit_breaker_for(&spec(
541                "e",
542                ModelTask::Embed,
543                "text-embedding-005",
544                serde_json::Value::Null,
545            )),
546            model_id: "text-embedding-005".into(),
547            options: opts,
548            dimensions: 768,
549        };
550        assert_eq!(
551            emb.endpoint_url(),
552            "https://aiplatform.googleapis.com/v1/projects/p/locations/global/publishers/google/models/text-embedding-005:predict"
553        );
554    }
555
556    #[tokio::test]
557    async fn breaker_reused_for_same_runtime_key() {
558        let _lock = ENV_LOCK.lock().await;
559        // SAFETY: protected by ENV_LOCK
560        unsafe {
561            std::env::set_var("VERTEX_AI_TOKEN", "test-token");
562            std::env::set_var("VERTEX_AI_PROJECT", "test-project");
563        }
564
565        let provider = RemoteVertexAIProvider::new();
566        let s1 = spec(
567            "embed/a",
568            ModelTask::Embed,
569            "text-embedding-005",
570            serde_json::Value::Null,
571        );
572        let s2 = spec(
573            "embed/b",
574            ModelTask::Embed,
575            "text-embedding-005",
576            serde_json::Value::Null,
577        );
578
579        let _ = provider.load(&s1).await.unwrap();
580        let _ = provider.load(&s2).await.unwrap();
581
582        assert_eq!(provider.breaker_count(), 1);
583
584        // SAFETY: protected by ENV_LOCK
585        unsafe {
586            std::env::remove_var("VERTEX_AI_TOKEN");
587            std::env::remove_var("VERTEX_AI_PROJECT");
588        }
589    }
590
591    #[tokio::test]
592    async fn breaker_cleanup_evicts_stale_entries() {
593        let _lock = ENV_LOCK.lock().await;
594        // SAFETY: protected by ENV_LOCK
595        unsafe {
596            std::env::set_var("VERTEX_AI_TOKEN", "test-token");
597            std::env::set_var("VERTEX_AI_PROJECT", "test-project");
598        }
599
600        let provider = RemoteVertexAIProvider::new();
601        let stale = spec(
602            "embed/stale",
603            ModelTask::Embed,
604            "text-embedding-005",
605            serde_json::Value::Null,
606        );
607        let fresh = spec(
608            "embed/fresh",
609            ModelTask::Embed,
610            "text-embedding-004",
611            serde_json::Value::Null,
612        );
613        provider.insert_test_breaker(
614            ModelRuntimeKey::new(&stale),
615            RemoteProviderBase::BREAKER_TTL + Duration::from_secs(5),
616        );
617        provider.insert_test_breaker(ModelRuntimeKey::new(&fresh), Duration::from_secs(1));
618        assert_eq!(provider.breaker_count(), 2);
619
620        provider.force_cleanup_now_for_test();
621        let _ = provider.load(&fresh).await.unwrap();
622        assert_eq!(provider.breaker_count(), 1);
623
624        // SAFETY: protected by ENV_LOCK
625        unsafe {
626            std::env::remove_var("VERTEX_AI_TOKEN");
627            std::env::remove_var("VERTEX_AI_PROJECT");
628        }
629    }
630
631    #[tokio::test]
632    async fn load_fails_when_project_is_missing() {
633        let _lock = ENV_LOCK.lock().await;
634        // SAFETY: protected by ENV_LOCK
635        unsafe {
636            std::env::set_var("VERTEX_AI_TOKEN", "test-token");
637            std::env::remove_var("VERTEX_AI_PROJECT");
638        }
639
640        let provider = RemoteVertexAIProvider::new();
641        let s = spec(
642            "embed/a",
643            ModelTask::Embed,
644            "text-embedding-005",
645            serde_json::Value::Null,
646        );
647        let err = provider.load(&s).await.unwrap_err();
648        assert!(err.to_string().contains("VERTEX_AI_PROJECT"));
649
650        // SAFETY: protected by ENV_LOCK
651        unsafe {
652            std::env::remove_var("VERTEX_AI_TOKEN");
653        }
654    }
655
656    #[test]
657    fn generation_payload_alternates_roles() {
658        use crate::traits::Message;
659        let messages = vec![
660            Message::user("user question"),
661            Message::assistant("assistant answer"),
662            Message::user("user follow-up"),
663        ];
664        let payload = build_google_generate_payload(&messages, &GenerationOptions::default());
665        let contents = payload["contents"].as_array().unwrap();
666
667        assert_eq!(contents[0]["role"], "user");
668        assert_eq!(contents[1]["role"], "model");
669        assert_eq!(contents[2]["role"], "user");
670    }
671
672    #[test]
673    fn generation_payload_includes_generation_options() {
674        use crate::traits::Message;
675        let messages = vec![Message::user("hello")];
676        let payload = build_google_generate_payload(
677            &messages,
678            &GenerationOptions {
679                max_tokens: Some(64),
680                temperature: Some(0.7),
681                top_p: Some(0.9),
682                ..Default::default()
683            },
684        );
685
686        assert_eq!(payload["generationConfig"]["maxOutputTokens"], 64);
687        let temperature = payload["generationConfig"]["temperature"].as_f64().unwrap();
688        let top_p = payload["generationConfig"]["topP"].as_f64().unwrap();
689        assert!((temperature - 0.7).abs() < 1e-6);
690        assert!((top_p - 0.9).abs() < 1e-6);
691    }
692}
uni_xervo/provider/vertexai.rs

uni_xervo/provider/
vertexai.rs
