rig_core/providers/cohere/
client.rs

1use crate::{
2    Embed,
3    client::{self, BearerAuth, DebugExt, Provider},
4    embeddings::EmbeddingsBuilder,
5    http_client::HttpClientExt,
6    wasm_compat::*,
7};
8
9use super::{CompletionModel, EmbeddingModel, ImageEmbeddingModel};
10use serde::Deserialize;
11
12// ================================================================
13// Main Cohere Client
14// ================================================================
15
16#[derive(Debug, Default, Clone, Copy)]
17pub struct CohereExt;
18
19#[derive(Debug, Default, Clone, Copy)]
20pub struct CohereBuilder;
21
22type CohereApiKey = BearerAuth;
23
24pub type Client<H = reqwest::Client> = client::Client<CohereExt, H>;
25pub type ClientBuilder<H = crate::markers::Missing> =
26    client::ClientBuilder<CohereBuilder, CohereApiKey, H>;
27
28impl Provider for CohereExt {
29    type Builder = CohereBuilder;
30    const VERIFY_PATH: &'static str = "/models";
31}
32
33client::impl_capabilities!(
34    CohereExt,
35    completion = CompletionModel<H>,
36    embeddings = EmbeddingModel<H>,
37);
38
39impl DebugExt for CohereExt {}
40
41client::impl_default_provider_builder!(
42    CohereBuilder => CohereExt,
43    api_key = CohereApiKey,
44    base_url = "https://api.cohere.ai",
45);
46
47client::impl_provider_client!(Client, input = CohereApiKey, api_key_env = "COHERE_API_KEY",);
48
49#[derive(Debug)]
50pub struct ApiErrorResponse {
51    /// Provider error message; tolerant of `{"message": "..."}`,
52    /// `{"error": "..."}`, nested `{"error": {"message": ...}}`, and bodies
53    /// carrying both keys. Used for logging only — the raw body is preserved
54    /// on the returned error.
55    pub message: String,
56}
57
58impl<'de> Deserialize<'de> for ApiErrorResponse {
59    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
60    where
61        D: serde::Deserializer<'de>,
62    {
63        Ok(Self {
64            message: crate::providers::internal::envelope::error_message(deserializer)?,
65        })
66    }
67}
68
69#[derive(Debug, Deserialize)]
70#[serde(untagged)]
71pub enum ApiResponse<T> {
72    Ok(T),
73    Err(ApiErrorResponse),
74}
75
76impl<T> Client<T>
77where
78    T: HttpClientExt + Clone + WasmCompatSend + WasmCompatSync + 'static,
79{
80    pub fn embeddings<D: Embed>(
81        &self,
82        model: impl Into<String>,
83        input_type: &str,
84    ) -> EmbeddingsBuilder<EmbeddingModel<T>, D> {
85        EmbeddingsBuilder::new(self.embedding_model(model, input_type))
86    }
87
88    /// Note: default embedding dimension of 0 will be used if model is not known.
89    /// If this is the case, it's better to use function `embedding_model_with_ndims`
90    pub fn embedding_model(&self, model: impl Into<String>, input_type: &str) -> EmbeddingModel<T> {
91        let model = model.into();
92        let ndims = super::model_dimensions_from_identifier(&model).unwrap_or_default();
93
94        EmbeddingModel::new(self.clone(), model, input_type, ndims)
95    }
96
97    /// Create a Cohere `embed-english-v3.0` model for embedding PNG, JPEG,
98    /// WebP, or GIF bytes.
99    ///
100    /// Images must be at least 2×2 pixels and no larger than 5 MB.
101    /// Cohere accepts one image per request, so
102    /// [`crate::embeddings::ImageEmbeddingModel::embed_images`] sends batches
103    /// as ordered individual requests.
104    pub fn image_embedding_model(&self) -> ImageEmbeddingModel<T> {
105        ImageEmbeddingModel::new(self.clone())
106    }
107
108    /// Create an embedding model with the given name and the number of dimensions in the embedding generated by the model.
109    pub fn embedding_model_with_ndims(
110        &self,
111        model: impl Into<String>,
112        input_type: &str,
113        ndims: usize,
114    ) -> EmbeddingModel<T> {
115        EmbeddingModel::new(self.clone(), model, input_type, ndims)
116    }
117}
118#[cfg(test)]
119mod tests {
120    #[test]
121    fn test_client_initialization() {
122        let _client =
123            crate::providers::cohere::Client::new("dummy-key").expect("Client::new() failed");
124        let _client_from_builder = crate::providers::cohere::Client::builder()
125            .api_key("dummy-key")
126            .build()
127            .expect("Client::builder() failed");
128    }
129}
rig_core/providers/cohere/client.rs

rig_core/providers/cohere/
client.rs
