use std::env;
use crate::ProviderError;
pub const COHERE_BASE_URL: &str = "https://api.cohere.com/v2";
pub const COHERE_MODELS: [&str; 4] = ["command-r-plus", "command-r", "command", "command-light"];
#[derive(Debug, Clone)]
pub struct CohereConfig {
pub api_key: String,
pub base_url: String,
pub model: String,
pub temperature: Option<f32>,
pub max_tokens: Option<usize>,
pub preamble: Option<String>,
}
impl Default for CohereConfig {
fn default() -> Self {
Self {
api_key: String::new(),
base_url: COHERE_BASE_URL.to_string(),
model: "command-r-plus".to_string(),
temperature: None,
max_tokens: None,
preamble: None,
}
}
}
impl CohereConfig {
pub fn new(api_key: impl Into<String>) -> Self {
Self {
api_key: api_key.into(),
..Default::default()
}
}
pub fn from_env_result() -> Result<Self, ProviderError> {
let api_key = env::var("COHERE_API_KEY").map_err(|_| {
ProviderError::Config("COHERE_API_KEY environment variable not set".to_string())
})?;
let base_url = env::var("COHERE_BASE_URL").unwrap_or_else(|_| COHERE_BASE_URL.to_string());
let model = env::var("COHERE_MODEL").unwrap_or_else(|_| "command-r-plus".to_string());
Ok(Self {
api_key,
base_url,
model,
..Default::default()
})
}
pub fn with_model(mut self, model: impl Into<String>) -> Self {
self.model = model.into();
self
}
pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
self.base_url = url.into();
self
}
pub fn with_temperature(mut self, temp: f32) -> Self {
self.temperature = Some(temp);
self
}
pub fn with_max_tokens(mut self, max: usize) -> Self {
self.max_tokens = Some(max);
self
}
pub fn with_preamble(mut self, preamble: impl Into<String>) -> Self {
self.preamble = Some(preamble.into());
self
}
}