Skip to main content

rig/providers/cohere/
mod.rs

1//! Cohere API client and Rig integration
2//!
3//! # Example
4//! ```
5//! use rig::providers::cohere;
6//!
7//! let client = cohere::Client::new("YOUR_API_KEY");
8//!
9//! let command_r = client.completion_model(cohere::COMMAND_R);
10//! ```
11
12pub mod client;
13pub mod completion;
14pub mod embeddings;
15pub mod streaming;
16
17pub use client::{ApiErrorResponse, ApiResponse, Client};
18pub use completion::CompletionModel;
19pub use embeddings::EmbeddingModel;
20
21// ================================================================
22// Cohere Completion Models
23// ================================================================
24
25/// `command-r-plus` completion model
26pub const COMMAND_R_PLUS: &str = "command-r-plus";
27/// `command-r` completion model
28pub const COMMAND_R: &str = "command-r";
29/// `command` completion model
30pub const COMMAND: &str = "command";
31/// `command-nightly` completion model
32pub const COMMAND_NIGHTLY: &str = "command-nightly";
33/// `command-light` completion model
34pub const COMMAND_LIGHT: &str = "command-light";
35/// `command-light-nightly` completion model
36pub const COMMAND_LIGHT_NIGHTLY: &str = "command-light-nightly";
37
38// ================================================================
39// Cohere Embedding Models
40// ================================================================
41
42/// `embed-english-v3.0` embedding model
43pub const EMBED_ENGLISH_V3: &str = "embed-english-v3.0";
44/// `embed-english-light-v3.0` embedding model
45pub const EMBED_ENGLISH_LIGHT_V3: &str = "embed-english-light-v3.0";
46/// `embed-multilingual-v3.0` embedding model
47pub const EMBED_MULTILINGUAL_V3: &str = "embed-multilingual-v3.0";
48/// `embed-multilingual-light-v3.0` embedding model
49pub const EMBED_MULTILINGUAL_LIGHT_V3: &str = "embed-multilingual-light-v3.0";
50
51pub(crate) fn model_dimensions_from_identifier(identifier: &str) -> Option<usize> {
52    match identifier {
53        EMBED_ENGLISH_V3 | EMBED_MULTILINGUAL_V3 => Some(1_024),
54        EMBED_ENGLISH_LIGHT_V3 | EMBED_MULTILINGUAL_LIGHT_V3 => Some(384),
55        _ => None,
56    }
57}