uni_xervo/provider/
cohere.rs

1use crate::api::{ModelAliasSpec, ModelTask};
2use crate::error::{Result, RuntimeError};
3use crate::provider::remote_common::{RemoteProviderBase, check_http_status, resolve_api_key};
4use crate::traits::{
5    EmbedResult, EmbeddingModel, GenerationOptions, GenerationResult, GeneratorModel, ImageInput,
6    LoadedModelHandle, Message, MessageRole, Modality, ModelProvider, MultimodalBlock,
7    MultimodalEmbeddingModel, MultimodalInput, ProviderCapabilities, ProviderHealth, RerankerModel,
8    ScoredDoc, TokenUsage,
9};
10use async_trait::async_trait;
11use base64::Engine;
12use reqwest::Client;
13use serde_json::json;
14use std::sync::Arc;
15
16/// Remote provider that calls the [Cohere API](https://docs.cohere.com/reference/about)
17/// for embedding, text generation (chat), and reranking.
18///
19/// Requires the `CO_API_KEY` environment variable (or a custom env var name
20/// via the `api_key_env` option).
21pub struct RemoteCohereProvider {
22    base: RemoteProviderBase,
23}
24
25impl Default for RemoteCohereProvider {
26    fn default() -> Self {
27        Self {
28            base: RemoteProviderBase::new(),
29        }
30    }
31}
32
33impl RemoteCohereProvider {
34    pub fn new() -> Self {
35        Self::default()
36    }
37
38    #[cfg(test)]
39    fn insert_test_breaker(&self, key: crate::api::ModelRuntimeKey, age: std::time::Duration) {
40        self.base.insert_test_breaker(key, age);
41    }
42
43    #[cfg(test)]
44    fn breaker_count(&self) -> usize {
45        self.base.breaker_count()
46    }
47
48    #[cfg(test)]
49    fn force_cleanup_now_for_test(&self) {
50        self.base.force_cleanup_now_for_test();
51    }
52}
53
54#[async_trait]
55impl ModelProvider for RemoteCohereProvider {
56    fn provider_id(&self) -> &'static str {
57        "remote/cohere"
58    }
59
60    fn capabilities(&self) -> ProviderCapabilities {
61        ProviderCapabilities {
62            supported_tasks: vec![
63                ModelTask::Embed,
64                ModelTask::Generate,
65                ModelTask::Rerank,
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", "CO_API_KEY")?;
74
75        let input_type = spec
76            .options
77            .get("input_type")
78            .and_then(|v| v.as_str())
79            .unwrap_or("search_document")
80            .to_string();
81        let embedding_dimensions = spec
82            .options
83            .get("embedding_dimensions")
84            .and_then(|v| v.as_u64())
85            .map(|v| v as u32);
86
87        match spec.task {
88            ModelTask::Embed => {
89                let default_dims = match spec.model_id.as_str() {
90                    "embed-english-light-v3.0" | "embed-multilingual-light-v3.0" => 384,
91                    "embed-v4.0" => 1536,
92                    _ => 1024,
93                };
94                let model = CohereEmbeddingModel {
95                    client: self.base.client.clone(),
96                    cb: cb.clone(),
97                    model_id: spec.model_id.clone(),
98                    api_key,
99                    input_type,
100                    dimensions: embedding_dimensions.unwrap_or(default_dims),
101                };
102                let handle: Arc<dyn EmbeddingModel> = Arc::new(model);
103                Ok(Arc::new(handle) as LoadedModelHandle)
104            }
105            ModelTask::Generate => {
106                let model = CohereGeneratorModel {
107                    client: self.base.client.clone(),
108                    cb,
109                    model_id: spec.model_id.clone(),
110                    api_key,
111                };
112                let handle: Arc<dyn GeneratorModel> = Arc::new(model);
113                Ok(Arc::new(handle) as LoadedModelHandle)
114            }
115            ModelTask::Rerank => {
116                let model = CohereRerankerModel {
117                    client: self.base.client.clone(),
118                    cb,
119                    model_id: spec.model_id.clone(),
120                    api_key,
121                };
122                let handle: Arc<dyn RerankerModel> = Arc::new(model);
123                Ok(Arc::new(handle) as LoadedModelHandle)
124            }
125            ModelTask::EmbedMultimodal => {
126                // Cohere Embed v4 returns 1536-dim for the canonical model
127                // and is the only multimodal Cohere embedder at time of
128                // writing. Other model_ids fall through to the same default
129                // and can be overridden via `embedding_dimensions`.
130                let default_dims = match spec.model_id.as_str() {
131                    "embed-v4.0" => 1536,
132                    _ => 1024,
133                };
134                let model = CohereMultimodalEmbeddingModel {
135                    client: self.base.client.clone(),
136                    cb,
137                    model_id: spec.model_id.clone(),
138                    api_key,
139                    input_type,
140                    dimensions: embedding_dimensions.unwrap_or(default_dims),
141                };
142                let handle: Arc<dyn MultimodalEmbeddingModel> = Arc::new(model);
143                Ok(Arc::new(handle) as LoadedModelHandle)
144            }
145            task => Err(RuntimeError::CapabilityMismatch(format!(
146                "Cohere provider does not support task {:?}",
147                task
148            ))),
149        }
150    }
151
152    async fn health(&self) -> ProviderHealth {
153        ProviderHealth::Healthy
154    }
155}
156
157struct CohereEmbeddingModel {
158    client: Client,
159    cb: crate::reliability::CircuitBreakerWrapper,
160    model_id: String,
161    api_key: String,
162    input_type: String,
163    dimensions: u32,
164}
165
166#[async_trait]
167impl EmbeddingModel for CohereEmbeddingModel {
168    async fn embed(&self, texts: &[&str]) -> Result<EmbedResult> {
169        let texts: Vec<String> = texts.iter().map(|s| s.to_string()).collect();
170
171        self.cb
172            .call(move || async move {
173                let response = self
174                    .client
175                    .post("https://api.cohere.com/v2/embed")
176                    .header("Authorization", format!("Bearer {}", self.api_key))
177                    .json(&json!({
178                        "texts": texts,
179                        "model": self.model_id,
180                        "input_type": self.input_type,
181                        "embedding_types": ["float"]
182                    }))
183                    .send()
184                    .await
185                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
186
187                let body: serde_json::Value = check_http_status("Cohere", response)?
188                    .json()
189                    .await
190                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
191
192                let float_embeddings = body
193                    .get("embeddings")
194                    .and_then(|e| e.get("float"))
195                    .and_then(|f| f.as_array())
196                    .ok_or_else(|| {
197                        RuntimeError::ApiError(
198                            "Invalid Cohere embedding response format".to_string(),
199                        )
200                    })?;
201
202                let mut vectors = Vec::new();
203                for embedding in float_embeddings {
204                    if let Some(values) = embedding.as_array() {
205                        let vec: Vec<f32> = values
206                            .iter()
207                            .filter_map(|v| v.as_f64().map(|f| f as f32))
208                            .collect();
209                        vectors.push(vec);
210                    }
211                }
212
213                // Cohere v2 embed response carries optional
214                // `meta.billed_units.input_tokens`.
215                let usage = body
216                    .get("meta")
217                    .and_then(|m| m.get("billed_units"))
218                    .and_then(|b| b.get("input_tokens"))
219                    .and_then(|t| t.as_u64())
220                    .map(|tokens| TokenUsage {
221                        prompt_tokens: tokens as usize,
222                        completion_tokens: 0,
223                        total_tokens: tokens as usize,
224                    });
225
226                Ok(EmbedResult { vectors, usage })
227            })
228            .await
229    }
230
231    fn dimensions(&self) -> u32 {
232        self.dimensions
233    }
234}
235
236impl crate::traits::ModelInfo for CohereEmbeddingModel {
237    fn model_id(&self) -> &str {
238        &self.model_id
239    }
240}
241
242struct CohereGeneratorModel {
243    client: Client,
244    cb: crate::reliability::CircuitBreakerWrapper,
245    model_id: String,
246    api_key: String,
247}
248
249impl crate::traits::ModelInfo for CohereGeneratorModel {
250    fn model_id(&self) -> &str {
251        &self.model_id
252    }
253}
254
255#[async_trait]
256impl GeneratorModel for CohereGeneratorModel {
257    async fn generate(
258        &self,
259        messages: &[Message],
260        options: GenerationOptions,
261    ) -> Result<GenerationResult> {
262        let messages: Vec<serde_json::Value> = messages
263            .iter()
264            .map(|msg| {
265                let role = match msg.role {
266                    MessageRole::System => "system",
267                    MessageRole::User => "user",
268                    MessageRole::Assistant => "assistant",
269                };
270                json!({ "role": role, "content": msg.text() })
271            })
272            .collect();
273
274        self.cb
275            .call(move || async move {
276                let mut body = json!({
277                    "model": self.model_id,
278                    "messages": messages,
279                });
280
281                if let Some(max_tokens) = options.max_tokens {
282                    body["max_tokens"] = json!(max_tokens);
283                }
284                if let Some(temperature) = options.temperature {
285                    body["temperature"] = json!(temperature);
286                }
287                if let Some(top_p) = options.top_p {
288                    body["p"] = json!(top_p);
289                }
290
291                let response = self
292                    .client
293                    .post("https://api.cohere.com/v2/chat")
294                    .header("Authorization", format!("Bearer {}", self.api_key))
295                    .json(&body)
296                    .send()
297                    .await
298                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
299
300                let body: serde_json::Value = check_http_status("Cohere", response)?
301                    .json()
302                    .await
303                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
304
305                let text = body
306                    .get("message")
307                    .and_then(|m| m.get("content"))
308                    .and_then(|c| c.as_array())
309                    .and_then(|arr| arr.first())
310                    .and_then(|item| item.get("text"))
311                    .and_then(|t| t.as_str())
312                    .unwrap_or("")
313                    .to_string();
314
315                let usage = body.get("usage").map(|u| {
316                    let input = u
317                        .get("tokens")
318                        .and_then(|t| t.get("input_tokens"))
319                        .and_then(|v| v.as_u64())
320                        .unwrap_or(0);
321                    let output = u
322                        .get("tokens")
323                        .and_then(|t| t.get("output_tokens"))
324                        .and_then(|v| v.as_u64())
325                        .unwrap_or(0);
326                    TokenUsage {
327                        prompt_tokens: input as usize,
328                        completion_tokens: output as usize,
329                        total_tokens: (input + output) as usize,
330                    }
331                });
332
333                Ok(GenerationResult {
334                    text,
335                    usage,
336                    images: vec![],
337                    audio: None,
338                })
339            })
340            .await
341    }
342}
343
344struct CohereRerankerModel {
345    client: Client,
346    cb: crate::reliability::CircuitBreakerWrapper,
347    model_id: String,
348    api_key: String,
349}
350
351impl crate::traits::ModelInfo for CohereRerankerModel {
352    fn model_id(&self) -> &str {
353        &self.model_id
354    }
355}
356
357#[async_trait]
358impl RerankerModel for CohereRerankerModel {
359    async fn rerank(&self, query: &str, docs: &[&str]) -> Result<Vec<ScoredDoc>> {
360        let query = query.to_string();
361        let docs: Vec<String> = docs.iter().map(|s| s.to_string()).collect();
362
363        self.cb
364            .call(move || async move {
365                let response = self
366                    .client
367                    .post("https://api.cohere.com/v2/rerank")
368                    .header("Authorization", format!("Bearer {}", self.api_key))
369                    .json(&json!({
370                        "query": query,
371                        "documents": docs,
372                        "model": self.model_id,
373                    }))
374                    .send()
375                    .await
376                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
377
378                let body: serde_json::Value = check_http_status("Cohere", response)?
379                    .json()
380                    .await
381                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
382
383                let results_json =
384                    body.get("results")
385                        .and_then(|r| r.as_array())
386                        .ok_or_else(|| {
387                            RuntimeError::ApiError("Invalid rerank response format".to_string())
388                        })?;
389
390                let mut results = Vec::new();
391                for item in results_json {
392                    let index = item.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
393                    let score = item
394                        .get("relevance_score")
395                        .and_then(|s| s.as_f64())
396                        .unwrap_or(0.0) as f32;
397                    results.push(ScoredDoc {
398                        index,
399                        score,
400                        text: None,
401                    });
402                }
403                Ok(results)
404            })
405            .await
406    }
407}
408
409struct CohereMultimodalEmbeddingModel {
410    client: Client,
411    cb: crate::reliability::CircuitBreakerWrapper,
412    model_id: String,
413    api_key: String,
414    input_type: String,
415    dimensions: u32,
416}
417
418impl CohereMultimodalEmbeddingModel {
419    /// Convert one [`MultimodalInput`] into Cohere v4's `{ content: [...] }`
420    /// shape. Text and image blocks are accepted; audio is rejected because
421    /// Cohere Embed v4 does not support audio input as of this writing.
422    fn input_to_json(input: &MultimodalInput) -> Result<serde_json::Value> {
423        let mut parts = Vec::with_capacity(input.blocks.len());
424        for block in &input.blocks {
425            match block {
426                MultimodalBlock::Text(text) => {
427                    parts.push(json!({ "type": "text", "text": text }));
428                }
429                MultimodalBlock::Image(ImageInput::Url(url)) => {
430                    parts.push(json!({
431                        "type": "image_url",
432                        "image_url": { "url": url }
433                    }));
434                }
435                MultimodalBlock::Image(ImageInput::Bytes { data, media_type }) => {
436                    let b64 = base64::engine::general_purpose::STANDARD.encode(data);
437                    let data_url = format!("data:{media_type};base64,{b64}");
438                    parts.push(json!({
439                        "type": "image_url",
440                        "image_url": { "url": data_url }
441                    }));
442                }
443                MultimodalBlock::Audio(_) => {
444                    return Err(RuntimeError::ApiError(
445                        "Cohere Embed v4 does not support audio inputs".to_string(),
446                    ));
447                }
448            }
449        }
450        Ok(json!({ "content": parts }))
451    }
452}
453
454#[async_trait]
455impl MultimodalEmbeddingModel for CohereMultimodalEmbeddingModel {
456    async fn embed(&self, inputs: Vec<MultimodalInput>) -> Result<EmbedResult> {
457        let body_inputs: Result<Vec<serde_json::Value>> =
458            inputs.iter().map(Self::input_to_json).collect();
459        let body_inputs = body_inputs?;
460        let model = self.model_id.clone();
461        let input_type = self.input_type.clone();
462        let api_key = self.api_key.clone();
463        let client = self.client.clone();
464
465        self.cb
466            .call(move || async move {
467                let response = client
468                    .post("https://api.cohere.com/v2/embed")
469                    .header("Authorization", format!("Bearer {api_key}"))
470                    .json(&json!({
471                        "model": model,
472                        "input_type": input_type,
473                        "embedding_types": ["float"],
474                        "inputs": body_inputs,
475                    }))
476                    .send()
477                    .await
478                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
479
480                let body: serde_json::Value = check_http_status("Cohere", response)?
481                    .json()
482                    .await
483                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
484
485                let float_embeddings = body
486                    .get("embeddings")
487                    .and_then(|e| e.get("float"))
488                    .and_then(|f| f.as_array())
489                    .ok_or_else(|| {
490                        RuntimeError::ApiError(
491                            "Invalid Cohere multimodal embedding response format".to_string(),
492                        )
493                    })?;
494
495                let mut vectors = Vec::with_capacity(float_embeddings.len());
496                for embedding in float_embeddings {
497                    if let Some(values) = embedding.as_array() {
498                        let vec: Vec<f32> = values
499                            .iter()
500                            .filter_map(|v| v.as_f64().map(|f| f as f32))
501                            .collect();
502                        vectors.push(vec);
503                    }
504                }
505
506                // Cohere v2 embed response carries optional `meta.billed_units.input_tokens`.
507                let usage = body
508                    .get("meta")
509                    .and_then(|m| m.get("billed_units"))
510                    .and_then(|b| b.get("input_tokens"))
511                    .and_then(|t| t.as_u64())
512                    .map(|tokens| TokenUsage {
513                        prompt_tokens: tokens as usize,
514                        completion_tokens: 0,
515                        total_tokens: tokens as usize,
516                    });
517
518                Ok(EmbedResult { vectors, usage })
519            })
520            .await
521    }
522
523    fn dimensions(&self) -> u32 {
524        self.dimensions
525    }
526
527    fn supported_modalities(&self) -> &[Modality] {
528        // Cohere Embed v4 supports text + image only (no audio, no video).
529        const COHERE_V4_MODALITIES: &[Modality] = &[Modality::Text, Modality::Image];
530        COHERE_V4_MODALITIES
531    }
532}
533
534impl crate::traits::ModelInfo for CohereMultimodalEmbeddingModel {
535    fn model_id(&self) -> &str {
536        &self.model_id
537    }
538}
539
540#[cfg(test)]
541mod tests {
542    use super::*;
543    use crate::api::ModelRuntimeKey;
544    use crate::provider::remote_common::RemoteProviderBase;
545    use crate::traits::ModelProvider;
546    use std::time::Duration;
547
548    static ENV_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
549
550    fn spec(alias: &str, task: ModelTask, model_id: &str) -> ModelAliasSpec {
551        ModelAliasSpec {
552            alias: alias.to_string(),
553            task,
554            provider_id: "remote/cohere".to_string(),
555            model_id: model_id.to_string(),
556            revision: None,
557            warmup: crate::api::WarmupPolicy::Lazy,
558            required: false,
559            timeout: None,
560            load_timeout: None,
561            retry: None,
562            options: serde_json::Value::Null,
563        }
564    }
565
566    #[tokio::test]
567    async fn breaker_reused_for_same_runtime_key() {
568        let _lock = ENV_LOCK.lock().await;
569        unsafe { std::env::set_var("CO_API_KEY", "test-key") };
570
571        let provider = RemoteCohereProvider::new();
572        let s1 = spec("embed/a", ModelTask::Embed, "embed-english-v3.0");
573        let s2 = spec("embed/b", ModelTask::Embed, "embed-english-v3.0");
574
575        let _ = provider.load(&s1).await.unwrap();
576        let _ = provider.load(&s2).await.unwrap();
577
578        assert_eq!(provider.breaker_count(), 1);
579
580        unsafe { std::env::remove_var("CO_API_KEY") };
581    }
582
583    #[tokio::test]
584    async fn breaker_isolated_by_task_and_model() {
585        let _lock = ENV_LOCK.lock().await;
586        unsafe { std::env::set_var("CO_API_KEY", "test-key") };
587
588        let provider = RemoteCohereProvider::new();
589        let embed = spec("embed/a", ModelTask::Embed, "embed-english-v3.0");
590        let gen_spec = spec("chat/a", ModelTask::Generate, "command-r-plus");
591        let rerank = spec("rerank/a", ModelTask::Rerank, "rerank-english-v3.0");
592
593        let _ = provider.load(&embed).await.unwrap();
594        let _ = provider.load(&gen_spec).await.unwrap();
595        let _ = provider.load(&rerank).await.unwrap();
596
597        assert_eq!(provider.breaker_count(), 3);
598
599        unsafe { std::env::remove_var("CO_API_KEY") };
600    }
601
602    #[tokio::test]
603    async fn breaker_cleanup_evicts_stale_entries() {
604        let _lock = ENV_LOCK.lock().await;
605        unsafe { std::env::set_var("CO_API_KEY", "test-key") };
606
607        let provider = RemoteCohereProvider::new();
608        let stale = spec("embed/stale", ModelTask::Embed, "embed-english-v3.0");
609        let fresh = spec("chat/fresh", ModelTask::Generate, "command-r-plus");
610        provider.insert_test_breaker(
611            ModelRuntimeKey::new(&stale),
612            RemoteProviderBase::BREAKER_TTL + Duration::from_secs(5),
613        );
614        provider.insert_test_breaker(ModelRuntimeKey::new(&fresh), Duration::from_secs(1));
615        assert_eq!(provider.breaker_count(), 2);
616
617        provider.force_cleanup_now_for_test();
618        let _ = provider.load(&fresh).await.unwrap();
619
620        assert_eq!(provider.breaker_count(), 1);
621
622        unsafe { std::env::remove_var("CO_API_KEY") };
623    }
624
625    #[tokio::test]
626    async fn supports_all_three_tasks() {
627        let _lock = ENV_LOCK.lock().await;
628        unsafe { std::env::set_var("CO_API_KEY", "test-key") };
629
630        let provider = RemoteCohereProvider::new();
631
632        let embed = spec("embed/a", ModelTask::Embed, "embed-english-v3.0");
633        assert!(provider.load(&embed).await.is_ok());
634
635        let gen_spec = spec("gen/a", ModelTask::Generate, "command-r-plus");
636        assert!(provider.load(&gen_spec).await.is_ok());
637
638        let rerank = spec("rerank/a", ModelTask::Rerank, "rerank-english-v3.0");
639        assert!(provider.load(&rerank).await.is_ok());
640
641        unsafe { std::env::remove_var("CO_API_KEY") };
642    }
643
644    #[tokio::test]
645    async fn rejects_raw_task() {
646        let _lock = ENV_LOCK.lock().await;
647        unsafe { std::env::set_var("CO_API_KEY", "test-key") };
648
649        let provider = RemoteCohereProvider::new();
650        let raw = spec("raw/a", ModelTask::Raw, "embed-english-v3.0");
651
652        let err = provider.load(&raw).await.unwrap_err();
653        assert!(matches!(err, RuntimeError::CapabilityMismatch(_)));
654        assert!(err.to_string().contains("does not support task Raw"));
655
656        unsafe { std::env::remove_var("CO_API_KEY") };
657    }
658
659    #[tokio::test]
660    async fn default_embedding_dimensions() {
661        let _lock = ENV_LOCK.lock().await;
662        unsafe { std::env::set_var("CO_API_KEY", "test-key") };
663
664        let provider = RemoteCohereProvider::new();
665        let s = spec("embed/dim", ModelTask::Embed, "embed-english-v3.0");
666
667        let handle = provider.load(&s).await.unwrap();
668        let model = handle.downcast_ref::<Arc<dyn EmbeddingModel>>().unwrap();
669        assert_eq!(model.dimensions(), 1024);
670
671        unsafe { std::env::remove_var("CO_API_KEY") };
672    }
673
674    #[tokio::test]
675    async fn custom_embedding_dimensions() {
676        let _lock = ENV_LOCK.lock().await;
677        unsafe { std::env::set_var("CO_API_KEY", "test-key") };
678
679        let provider = RemoteCohereProvider::new();
680        let mut s = spec("embed/dim-custom", ModelTask::Embed, "embed-english-v3.0");
681        s.options = serde_json::json!({"embedding_dimensions": 256});
682
683        let handle = provider.load(&s).await.unwrap();
684        let model = handle.downcast_ref::<Arc<dyn EmbeddingModel>>().unwrap();
685        assert_eq!(model.dimensions(), 256);
686
687        unsafe { std::env::remove_var("CO_API_KEY") };
688    }
689}
uni_xervo/provider/cohere.rs

uni_xervo/provider/
cohere.rs
