uni_xervo/provider/
gemini.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, resolve_api_key,
5};
6use crate::traits::{
7    AudioInput, EmbedResult, EmbeddingModel, GenerationOptions, GenerationResult, GeneratorModel,
8    ImageInput, LoadedModelHandle, Message, Modality, ModelProvider, MultimodalBlock,
9    MultimodalEmbeddingModel, MultimodalInput, ProviderCapabilities, ProviderHealth, TokenUsage,
10};
11use async_trait::async_trait;
12use base64::Engine;
13use reqwest::Client;
14use serde_json::json;
15use std::sync::Arc;
16
17/// Remote provider that calls the [Google Gemini API](https://ai.google.dev/api)
18/// for embedding (`batchEmbedContents`) and text generation (`generateContent`).
19///
20/// Requires the `GEMINI_API_KEY` environment variable (or a custom env var name
21/// via the `api_key_env` option).
22pub struct RemoteGeminiProvider {
23    base: RemoteProviderBase,
24}
25
26impl RemoteGeminiProvider {
27    pub fn new() -> Self {
28        Self::default()
29    }
30
31    #[cfg(test)]
32    fn insert_test_breaker(&self, key: crate::api::ModelRuntimeKey, age: std::time::Duration) {
33        self.base.insert_test_breaker(key, age);
34    }
35
36    #[cfg(test)]
37    fn breaker_count(&self) -> usize {
38        self.base.breaker_count()
39    }
40
41    #[cfg(test)]
42    fn force_cleanup_now_for_test(&self) {
43        self.base.force_cleanup_now_for_test();
44    }
45}
46
47impl Default for RemoteGeminiProvider {
48    fn default() -> Self {
49        Self {
50            base: RemoteProviderBase::new(),
51        }
52    }
53}
54
55#[async_trait]
56impl ModelProvider for RemoteGeminiProvider {
57    fn provider_id(&self) -> &'static str {
58        "remote/gemini"
59    }
60
61    fn capabilities(&self) -> ProviderCapabilities {
62        ProviderCapabilities {
63            supported_tasks: vec![
64                ModelTask::Embed,
65                ModelTask::Generate,
66                ModelTask::EmbedMultimodal,
67            ],
68        }
69    }
70
71    async fn load(&self, spec: &ModelAliasSpec) -> Result<LoadedModelHandle> {
72        let cb = self.base.circuit_breaker_for(spec);
73        let api_key = resolve_api_key(&spec.options, "api_key_env", "GEMINI_API_KEY")?;
74        let api_version = spec
75            .options
76            .get("api_version")
77            .and_then(|v| v.as_str())
78            .unwrap_or("v1beta")
79            .to_string();
80        let embedding_dimensions = spec
81            .options
82            .get("embedding_dimensions")
83            .and_then(|v| v.as_u64())
84            .map(|v| v as u32);
85
86        match spec.task {
87            ModelTask::Embed => {
88                let default_dims = match spec.model_id.as_str() {
89                    "gemini-embedding-001" | "gemini-embedding-exp-03-07" => 3072,
90                    _ => 768,
91                };
92                let model = GeminiEmbeddingModel {
93                    client: self.base.client.clone(),
94                    cb: cb.clone(),
95                    model_id: spec.model_id.clone(),
96                    api_key,
97                    api_version: api_version.clone(),
98                    dimensions: embedding_dimensions.unwrap_or(default_dims),
99                };
100                let handle: Arc<dyn EmbeddingModel> = Arc::new(model);
101                Ok(Arc::new(handle) as LoadedModelHandle)
102            }
103            ModelTask::Generate => {
104                let model = GeminiGeneratorModel {
105                    client: self.base.client.clone(),
106                    cb,
107                    model_id: spec.model_id.clone(),
108                    api_key,
109                    api_version,
110                };
111                let handle: Arc<dyn GeneratorModel> = Arc::new(model);
112                Ok(Arc::new(handle) as LoadedModelHandle)
113            }
114            ModelTask::EmbedMultimodal => {
115                // Gemini Embedding 2 returns 3072-dim by default.
116                let default_dims = 3072;
117                let model = GeminiMultimodalEmbeddingModel {
118                    client: self.base.client.clone(),
119                    cb,
120                    model_id: spec.model_id.clone(),
121                    api_key,
122                    api_version,
123                    dimensions: embedding_dimensions.unwrap_or(default_dims),
124                };
125                let handle: Arc<dyn MultimodalEmbeddingModel> = Arc::new(model);
126                Ok(Arc::new(handle) as LoadedModelHandle)
127            }
128            ModelTask::Raw => Err(RuntimeError::CapabilityMismatch(
129                "Gemini provider does not support task Raw".to_string(),
130            )),
131            _ => Err(RuntimeError::CapabilityMismatch(format!(
132                "Gemini provider does not support task {:?}",
133                spec.task
134            ))),
135        }
136    }
137
138    async fn health(&self) -> ProviderHealth {
139        ProviderHealth::Healthy
140    }
141}
142
143/// Embedding model backed by the Gemini batch embedding API.
144pub struct GeminiEmbeddingModel {
145    client: Client,
146    cb: crate::reliability::CircuitBreakerWrapper,
147    model_id: String,
148    api_key: String,
149    api_version: String,
150    dimensions: u32,
151}
152
153#[async_trait]
154impl EmbeddingModel for GeminiEmbeddingModel {
155    async fn embed(&self, texts: &[&str]) -> Result<EmbedResult> {
156        let texts: Vec<String> = texts.iter().map(|s| s.to_string()).collect();
157
158        self.cb
159            .call(move || async move {
160                let url = format!(
161                    "https://generativelanguage.googleapis.com/{}/models/{}:batchEmbedContents?key={}",
162                    self.api_version, self.model_id, self.api_key
163                );
164
165                let requests: Vec<_> = texts
166                    .iter()
167                    .map(|t| {
168                        json!({
169                            "model": format!("models/{}", self.model_id),
170                            "content": { "parts": [{ "text": t }] }
171                        })
172                    })
173                    .collect();
174
175                let response = self
176                    .client
177                    .post(&url)
178                    .json(&json!({ "requests": requests }))
179                    .send()
180                    .await
181                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
182
183                let body: serde_json::Value = check_http_status("Gemini", response)?
184                    .json()
185                    .await
186                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
187
188                let embeddings_json = body
189                    .get("embeddings")
190                    .and_then(|v| v.as_array())
191                    .ok_or_else(|| {
192                        RuntimeError::ApiError("Invalid response format".to_string())
193                    })?;
194
195                let mut vectors = Vec::new();
196                for item in embeddings_json {
197                    let values = item
198                        .get("values")
199                        .and_then(|v| v.as_array())
200                        .ok_or_else(|| {
201                            RuntimeError::ApiError("Missing values in embedding".to_string())
202                        })?;
203
204                    let vec: Vec<f32> = values
205                        .iter()
206                        .filter_map(|v| v.as_f64().map(|f| f as f32))
207                        .collect();
208                    vectors.push(vec);
209                }
210
211                // Gemini does not report usage on the batchEmbedContents endpoint.
212                Ok(EmbedResult {
213                    vectors,
214                    usage: None,
215                })
216            })
217            .await
218    }
219
220    fn dimensions(&self) -> u32 {
221        self.dimensions
222    }
223}
224
225impl crate::traits::ModelInfo for GeminiEmbeddingModel {
226    fn model_id(&self) -> &str {
227        &self.model_id
228    }
229}
230
231/// Text generation model backed by the Gemini `generateContent` API.
232pub struct GeminiGeneratorModel {
233    client: Client,
234    cb: crate::reliability::CircuitBreakerWrapper,
235    model_id: String,
236    api_key: String,
237    api_version: String,
238}
239
240impl crate::traits::ModelInfo for GeminiGeneratorModel {
241    fn model_id(&self) -> &str {
242        &self.model_id
243    }
244}
245
246#[async_trait]
247impl GeneratorModel for GeminiGeneratorModel {
248    async fn generate(
249        &self,
250        messages: &[Message],
251        options: GenerationOptions,
252    ) -> Result<GenerationResult> {
253        let messages: Vec<Message> = messages.to_vec();
254
255        self.cb
256            .call(move || async move {
257                let url = format!(
258                    "https://generativelanguage.googleapis.com/{}/models/{}:generateContent?key={}",
259                    self.api_version, self.model_id, self.api_key
260                );
261
262                let payload = build_google_generate_payload(&messages, &options);
263
264                let response = self
265                    .client
266                    .post(&url)
267                    .json(&payload)
268                    .send()
269                    .await
270                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
271
272                let body: serde_json::Value = check_http_status("Gemini", response)?
273                    .json()
274                    .await
275                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
276
277                let candidates = body
278                    .get("candidates")
279                    .and_then(|v| v.as_array())
280                    .ok_or_else(|| RuntimeError::ApiError("No candidates returned".to_string()))?;
281
282                let first_candidate = candidates
283                    .first()
284                    .ok_or_else(|| RuntimeError::ApiError("Empty candidates".to_string()))?;
285
286                let content_parts = first_candidate
287                    .get("content")
288                    .and_then(|c| c.get("parts"))
289                    .and_then(|p| p.as_array())
290                    .ok_or_else(|| RuntimeError::ApiError("Invalid content format".to_string()))?;
291
292                let text = content_parts
293                    .first()
294                    .and_then(|p| p.get("text"))
295                    .and_then(|t| t.as_str())
296                    .unwrap_or("")
297                    .to_string();
298
299                let usage = body.get("usageMetadata").map(|u| TokenUsage {
300                    prompt_tokens: u["promptTokenCount"].as_u64().unwrap_or(0) as usize,
301                    completion_tokens: u["candidatesTokenCount"].as_u64().unwrap_or(0) as usize,
302                    total_tokens: u["totalTokenCount"].as_u64().unwrap_or(0) as usize,
303                });
304
305                Ok(GenerationResult {
306                    text,
307                    usage,
308                    images: vec![],
309                    audio: None,
310                })
311            })
312            .await
313    }
314}
315
316/// Multimodal embedding model backed by Gemini Embedding 2's
317/// `batchEmbedContents` endpoint with multimodal `parts`.
318pub struct GeminiMultimodalEmbeddingModel {
319    client: Client,
320    cb: crate::reliability::CircuitBreakerWrapper,
321    model_id: String,
322    api_key: String,
323    api_version: String,
324    dimensions: u32,
325}
326
327impl GeminiMultimodalEmbeddingModel {
328    /// Convert one [`MultimodalInput`] into a Gemini `content.parts` array.
329    fn input_to_parts(input: &MultimodalInput) -> Vec<serde_json::Value> {
330        input
331            .blocks
332            .iter()
333            .map(|block| match block {
334                MultimodalBlock::Text(text) => json!({ "text": text }),
335                MultimodalBlock::Image(ImageInput::Url(url)) => json!({
336                    "file_data": { "file_uri": url }
337                }),
338                MultimodalBlock::Image(ImageInput::Bytes { data, media_type }) => {
339                    let b64 = base64::engine::general_purpose::STANDARD.encode(data);
340                    json!({
341                        "inline_data": { "mime_type": media_type, "data": b64 }
342                    })
343                }
344                MultimodalBlock::Audio(AudioInput::Bytes { data, media_type }) => {
345                    let b64 = base64::engine::general_purpose::STANDARD.encode(data);
346                    json!({
347                        "inline_data": { "mime_type": media_type, "data": b64 }
348                    })
349                }
350                MultimodalBlock::Audio(AudioInput::Pcm {
351                    sample_rate,
352                    samples,
353                    ..
354                }) => {
355                    // Build a minimal 16-bit PCM WAV in memory so the API
356                    // receives a self-describing container. Mono assumption
357                    // matches the typical embedding pipeline; multi-channel
358                    // callers should encode upstream and use Bytes.
359                    let wav_bytes = encode_pcm_to_wav(*sample_rate, samples);
360                    let b64 = base64::engine::general_purpose::STANDARD.encode(&wav_bytes);
361                    json!({
362                        "inline_data": { "mime_type": "audio/wav", "data": b64 }
363                    })
364                }
365            })
366            .collect()
367    }
368}
369
370/// Encode a slice of f32 PCM samples into a 16-bit mono WAV byte buffer.
371///
372/// Kept inline because Gemini is currently the only consumer; a shared
373/// audio-encode helper can land alongside PR-4 (whisper-cpp) which needs
374/// the inverse operation.
375fn encode_pcm_to_wav(sample_rate: u32, samples: &[f32]) -> Vec<u8> {
376    let bits_per_sample: u16 = 16;
377    let channels: u16 = 1;
378    let byte_rate = sample_rate * (bits_per_sample as u32) * (channels as u32) / 8;
379    let block_align = channels * bits_per_sample / 8;
380    let data_size = (samples.len() as u32) * (bits_per_sample as u32) * (channels as u32) / 8;
381
382    let mut buf = Vec::with_capacity(44 + data_size as usize);
383    buf.extend_from_slice(b"RIFF");
384    buf.extend_from_slice(&(36 + data_size).to_le_bytes());
385    buf.extend_from_slice(b"WAVE");
386    buf.extend_from_slice(b"fmt ");
387    buf.extend_from_slice(&16u32.to_le_bytes()); // fmt chunk size
388    buf.extend_from_slice(&1u16.to_le_bytes()); // PCM format
389    buf.extend_from_slice(&channels.to_le_bytes());
390    buf.extend_from_slice(&sample_rate.to_le_bytes());
391    buf.extend_from_slice(&byte_rate.to_le_bytes());
392    buf.extend_from_slice(&block_align.to_le_bytes());
393    buf.extend_from_slice(&bits_per_sample.to_le_bytes());
394    buf.extend_from_slice(b"data");
395    buf.extend_from_slice(&data_size.to_le_bytes());
396    for &sample in samples {
397        let clamped = sample.clamp(-1.0, 1.0);
398        let s = (clamped * i16::MAX as f32) as i16;
399        buf.extend_from_slice(&s.to_le_bytes());
400    }
401    buf
402}
403
404#[async_trait]
405impl MultimodalEmbeddingModel for GeminiMultimodalEmbeddingModel {
406    async fn embed(&self, inputs: Vec<MultimodalInput>) -> Result<EmbedResult> {
407        let requests: Vec<serde_json::Value> = inputs
408            .iter()
409            .map(|input| {
410                json!({
411                    "model": format!("models/{}", self.model_id),
412                    "content": { "parts": Self::input_to_parts(input) }
413                })
414            })
415            .collect();
416
417        self.cb
418            .call(move || async move {
419                let url = format!(
420                    "https://generativelanguage.googleapis.com/{}/models/{}:batchEmbedContents?key={}",
421                    self.api_version, self.model_id, self.api_key
422                );
423                let response = self
424                    .client
425                    .post(&url)
426                    .json(&json!({ "requests": requests }))
427                    .send()
428                    .await
429                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
430                let body: serde_json::Value = check_http_status("Gemini", response)?
431                    .json()
432                    .await
433                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
434
435                let embeddings_json = body
436                    .get("embeddings")
437                    .and_then(|v| v.as_array())
438                    .ok_or_else(|| {
439                        RuntimeError::ApiError(
440                            "Invalid Gemini multimodal embedding response format".to_string(),
441                        )
442                    })?;
443
444                let mut vectors = Vec::with_capacity(embeddings_json.len());
445                for item in embeddings_json {
446                    let values = item
447                        .get("values")
448                        .and_then(|v| v.as_array())
449                        .ok_or_else(|| {
450                            RuntimeError::ApiError("Missing values in embedding".to_string())
451                        })?;
452                    let vec: Vec<f32> = values
453                        .iter()
454                        .filter_map(|v| v.as_f64().map(|f| f as f32))
455                        .collect();
456                    vectors.push(vec);
457                }
458
459                // Gemini does not report usage on the embed endpoints.
460                Ok(EmbedResult {
461                    vectors,
462                    usage: None,
463                })
464            })
465            .await
466    }
467
468    fn dimensions(&self) -> u32 {
469        self.dimensions
470    }
471
472    fn supported_modalities(&self) -> &[Modality] {
473        const GEMINI_2_MODALITIES: &[Modality] = &[
474            Modality::Text,
475            Modality::Image,
476            Modality::Audio,
477            Modality::Video,
478        ];
479        GEMINI_2_MODALITIES
480    }
481}
482
483impl crate::traits::ModelInfo for GeminiMultimodalEmbeddingModel {
484    fn model_id(&self) -> &str {
485        &self.model_id
486    }
487}
488
489#[cfg(test)]
490mod tests {
491    use super::*;
492    use crate::api::ModelRuntimeKey;
493    use crate::provider::remote_common::RemoteProviderBase;
494    use crate::traits::ModelProvider;
495    use std::time::Duration;
496
497    static ENV_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
498
499    fn spec(alias: &str, task: ModelTask, model_id: &str) -> ModelAliasSpec {
500        ModelAliasSpec {
501            alias: alias.to_string(),
502            task,
503            provider_id: "remote/gemini".to_string(),
504            model_id: model_id.to_string(),
505            revision: None,
506            warmup: crate::api::WarmupPolicy::Lazy,
507            required: false,
508            timeout: None,
509            load_timeout: None,
510            retry: None,
511            options: serde_json::Value::Null,
512        }
513    }
514
515    #[tokio::test]
516    async fn breaker_reused_for_same_runtime_key() {
517        let _lock = ENV_LOCK.lock().await;
518        // SAFETY: protected by ENV_LOCK
519        unsafe { std::env::set_var("GEMINI_API_KEY", "test-key") };
520
521        let provider = RemoteGeminiProvider::new();
522        let s1 = spec("embed/a", ModelTask::Embed, "embedding-001");
523        let s2 = spec("embed/b", ModelTask::Embed, "embedding-001");
524
525        let _ = provider.load(&s1).await.unwrap();
526        let _ = provider.load(&s2).await.unwrap();
527
528        assert_eq!(provider.breaker_count(), 1);
529
530        // SAFETY: protected by ENV_LOCK
531        unsafe { std::env::remove_var("GEMINI_API_KEY") };
532    }
533
534    #[tokio::test]
535    async fn breaker_isolated_by_task_and_model() {
536        let _lock = ENV_LOCK.lock().await;
537        // SAFETY: protected by ENV_LOCK
538        unsafe { std::env::set_var("GEMINI_API_KEY", "test-key") };
539
540        let provider = RemoteGeminiProvider::new();
541        let embed = spec("embed/a", ModelTask::Embed, "embedding-001");
542        let gen_spec = spec("chat/a", ModelTask::Generate, "gemini-pro");
543
544        let _ = provider.load(&embed).await.unwrap();
545        let _ = provider.load(&gen_spec).await.unwrap();
546
547        assert_eq!(provider.breaker_count(), 2);
548
549        // SAFETY: protected by ENV_LOCK
550        unsafe { std::env::remove_var("GEMINI_API_KEY") };
551    }
552
553    #[tokio::test]
554    async fn breaker_cleanup_evicts_stale_entries() {
555        let _lock = ENV_LOCK.lock().await;
556        // SAFETY: protected by ENV_LOCK
557        unsafe { std::env::set_var("GEMINI_API_KEY", "test-key") };
558
559        let provider = RemoteGeminiProvider::new();
560        let stale = spec("embed/stale", ModelTask::Embed, "embedding-001");
561        let fresh = spec("embed/fresh", ModelTask::Embed, "embedding-002");
562        provider.insert_test_breaker(
563            ModelRuntimeKey::new(&stale),
564            RemoteProviderBase::BREAKER_TTL + Duration::from_secs(5),
565        );
566        provider.insert_test_breaker(ModelRuntimeKey::new(&fresh), Duration::from_secs(1));
567        assert_eq!(provider.breaker_count(), 2);
568
569        provider.force_cleanup_now_for_test();
570        let _ = provider.load(&fresh).await.unwrap();
571
572        assert_eq!(provider.breaker_count(), 1);
573
574        // SAFETY: protected by ENV_LOCK
575        unsafe { std::env::remove_var("GEMINI_API_KEY") };
576    }
577
578    #[test]
579    fn generation_payload_alternates_roles() {
580        use crate::traits::Message;
581        let messages = vec![
582            Message::user("user question"),
583            Message::assistant("assistant answer"),
584            Message::user("user follow-up"),
585        ];
586        let payload = build_google_generate_payload(&messages, &GenerationOptions::default());
587        let contents = payload["contents"].as_array().unwrap();
588
589        assert_eq!(contents[0]["role"], "user");
590        assert_eq!(contents[1]["role"], "model");
591        assert_eq!(contents[2]["role"], "user");
592    }
593
594    #[test]
595    fn generation_payload_includes_generation_options() {
596        use crate::traits::Message;
597        let messages = vec![Message::user("hello")];
598        let payload = build_google_generate_payload(
599            &messages,
600            &GenerationOptions {
601                max_tokens: Some(64),
602                temperature: Some(0.7),
603                top_p: Some(0.9),
604                ..Default::default()
605            },
606        );
607
608        assert_eq!(payload["generationConfig"]["maxOutputTokens"], 64);
609        let temperature = payload["generationConfig"]["temperature"].as_f64().unwrap();
610        let top_p = payload["generationConfig"]["topP"].as_f64().unwrap();
611        assert!((temperature - 0.7).abs() < 1e-6);
612        assert!((top_p - 0.9).abs() < 1e-6);
613    }
614
615    #[test]
616    fn generation_payload_extracts_system_instruction() {
617        use crate::traits::Message;
618        let messages = vec![Message::system("you are helpful"), Message::user("hello")];
619        let payload = build_google_generate_payload(&messages, &GenerationOptions::default());
620
621        // System message should be extracted into system_instruction
622        let si = &payload["system_instruction"];
623        assert_eq!(si["parts"][0]["text"], "you are helpful");
624
625        // Contents should only have the user message
626        let contents = payload["contents"].as_array().unwrap();
627        assert_eq!(contents.len(), 1);
628        assert_eq!(contents[0]["role"], "user");
629    }
630
631    #[test]
632    fn generation_payload_no_system_instruction_without_system_messages() {
633        use crate::traits::Message;
634        let messages = vec![Message::user("hello"), Message::assistant("hi")];
635        let payload = build_google_generate_payload(&messages, &GenerationOptions::default());
636
637        // No system_instruction field should be present
638        assert!(payload.get("system_instruction").is_none());
639
640        let contents = payload["contents"].as_array().unwrap();
641        assert_eq!(contents.len(), 2);
642    }
643
644    fn spec_with_opts(
645        alias: &str,
646        task: ModelTask,
647        model_id: &str,
648        options: serde_json::Value,
649    ) -> ModelAliasSpec {
650        ModelAliasSpec {
651            alias: alias.to_string(),
652            task,
653            provider_id: "remote/gemini".to_string(),
654            model_id: model_id.to_string(),
655            revision: None,
656            warmup: crate::api::WarmupPolicy::Lazy,
657            required: false,
658            timeout: None,
659            load_timeout: None,
660            retry: None,
661            options,
662        }
663    }
664
665    #[tokio::test]
666    async fn default_embedding_dimensions() {
667        let _lock = ENV_LOCK.lock().await;
668        // SAFETY: protected by ENV_LOCK
669        unsafe { std::env::set_var("GEMINI_API_KEY", "test-key") };
670
671        let provider = RemoteGeminiProvider::new();
672        let s = spec("embed/dim", ModelTask::Embed, "embedding-001");
673        let handle = provider.load(&s).await.unwrap();
674        let model = handle.downcast_ref::<Arc<dyn EmbeddingModel>>().unwrap();
675        assert_eq!(model.dimensions(), 768);
676
677        // SAFETY: protected by ENV_LOCK
678        unsafe { std::env::remove_var("GEMINI_API_KEY") };
679    }
680
681    #[tokio::test]
682    async fn custom_embedding_dimensions() {
683        let _lock = ENV_LOCK.lock().await;
684        // SAFETY: protected by ENV_LOCK
685        unsafe { std::env::set_var("GEMINI_API_KEY", "test-key") };
686
687        let provider = RemoteGeminiProvider::new();
688        let s = spec_with_opts(
689            "embed/dim-custom",
690            ModelTask::Embed,
691            "embedding-001",
692            json!({ "embedding_dimensions": 256 }),
693        );
694        let handle = provider.load(&s).await.unwrap();
695        let model = handle.downcast_ref::<Arc<dyn EmbeddingModel>>().unwrap();
696        assert_eq!(model.dimensions(), 256);
697
698        // SAFETY: protected by ENV_LOCK
699        unsafe { std::env::remove_var("GEMINI_API_KEY") };
700    }
701
702    #[tokio::test]
703    async fn api_version_option_accepted() {
704        let _lock = ENV_LOCK.lock().await;
705        // SAFETY: protected by ENV_LOCK
706        unsafe { std::env::set_var("GEMINI_API_KEY", "test-key") };
707
708        let provider = RemoteGeminiProvider::new();
709        let s = spec_with_opts(
710            "embed/v1",
711            ModelTask::Embed,
712            "embedding-001",
713            json!({ "api_version": "v1" }),
714        );
715        let handle = provider.load(&s).await;
716        assert!(handle.is_ok());
717
718        // SAFETY: protected by ENV_LOCK
719        unsafe { std::env::remove_var("GEMINI_API_KEY") };
720    }
721}
uni_xervo/provider/gemini.rs

uni_xervo/provider/
gemini.rs
