use secrecy::SecretString;
#[derive(Debug, Clone)]
pub struct GeneratorInfo {
pub name: String,
pub base_url: String,
pub model: String,
pub api_key: Option<SecretString>,
pub organization_id: Option<String>,
pub custom_headers: Vec<(String, String)>,
pub supports_streaming: bool,
pub supports_vision: bool,
pub supports_audio: bool,
pub max_context_length: Option<usize>,
pub default_params: super::CompletionParameters,
}
impl GeneratorInfo {
pub fn new(
name: impl Into<String>,
base_url: impl Into<String>,
model: impl Into<String>,
) -> Self {
Self {
name: name.into(),
base_url: base_url.into(),
model: model.into(),
api_key: None,
organization_id: None,
custom_headers: Vec::new(),
supports_streaming: true,
supports_vision: false,
supports_audio: false,
max_context_length: None,
default_params: super::CompletionParameters::default(),
}
}
pub fn with_api_key(mut self, key: impl Into<String>) -> Self {
self.api_key = Some(SecretString::from(key.into()));
self
}
pub fn with_api_key_from_env(mut self, env_var: &str) -> Self {
if let Ok(key) = std::env::var(env_var) {
self.api_key = Some(SecretString::from(key));
}
self
}
pub fn with_header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
self.custom_headers.push((name.into(), value.into()));
self
}
pub fn with_vision(mut self) -> Self {
self.supports_vision = true;
self
}
pub fn with_audio(mut self) -> Self {
self.supports_audio = true;
self
}
pub fn with_max_context(mut self, length: usize) -> Self {
self.max_context_length = Some(length);
self
}
pub fn with_default_params(mut self, params: super::CompletionParameters) -> Self {
self.default_params = params;
self
}
pub fn completions_url(&self) -> String {
let base = self.base_url.trim_end_matches('/');
format!("{}/chat/completions", base)
}
}
impl GeneratorInfo {
pub fn openrouter(model: impl Into<String>) -> Self {
Self::new("OpenRouter", "https://openrouter.ai/api/v1", model)
.with_api_key_from_env("OPENROUTER_API_KEY")
.with_header("HTTP-Referer", "https://github.com/minillmlib")
.with_header("X-Title", "MiniLLMLib")
}
pub fn openai(model: impl Into<String>) -> Self {
Self::new("OpenAI", "https://api.openai.com/v1", model)
.with_api_key_from_env("OPENAI_API_KEY")
}
pub fn anthropic(model: impl Into<String>) -> Self {
Self::new("Anthropic", "https://api.anthropic.com/v1", model)
.with_api_key_from_env("ANTHROPIC_API_KEY")
}
pub fn custom(
name: impl Into<String>,
base_url: impl Into<String>,
model: impl Into<String>,
) -> Self {
Self::new(name, base_url, model)
}
}