use std::borrow::Cow;
use std::fmt;
use std::hash::{Hash, Hasher};
use serde::de;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use crate::catalog::{Catalog, ModelFacts, ModelSpec, NotFound, Reference};
use crate::client::env::{self, EnvError};
use crate::driver::DynModel;
use crate::http_client::{DynHttpClient, HttpClientExt};
use crate::operation::Completion;
use crate::providers::{anthropic, gemini, openai};
use crate::serve::ErasedHandler;
use crate::serve::adapters::ModelAdapter;
use crate::wire::Secret;
pub(crate) const OPENAI_DIALECTS: &[&openai::wire::Dialect] = &[
&openai::wire::OPENAI,
&openai::wire::AZURE,
&openai::wire::DEEPSEEK,
&openai::wire::GROQ,
&openai::wire::HYPERBOLIC,
&openai::wire::MIRA,
&openai::wire::PERPLEXITY,
&openai::wire::TOGETHER,
&openai::wire::HUGGINGFACE,
&openai::wire::LLAMACPP,
&openai::wire::MISTRAL,
&openai::wire::OPENROUTER,
&openai::wire::VENICE,
&openai::wire::DOUBLEWORD,
&openai::wire::ZAI,
&openai::wire::MINIMAX,
&openai::wire::MOONSHOT,
&openai::wire::XIAOMIMIMO,
&openai::wire::COHERE,
&openai::wire::OLLAMA,
&crate::providers::xai::DIALECT,
&crate::providers::chatgpt::DIALECT,
&crate::providers::copilot::wire::DIALECT,
];
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum Format {
OpenAi,
Anthropic,
Gemini,
}
impl Format {
pub const ALL: [Format; 3] = [Format::OpenAi, Format::Anthropic, Format::Gemini];
pub fn as_str(self) -> &'static str {
match self {
Format::OpenAi => "openai",
Format::Anthropic => "anthropic",
Format::Gemini => "gemini",
}
}
pub fn named(name: &str) -> Option<Self> {
Self::ALL.into_iter().find(|format| format.as_str() == name)
}
fn names() -> String {
Self::ALL
.iter()
.map(|format| format.as_str())
.collect::<Vec<_>>()
.join(", ")
}
}
impl fmt::Display for Format {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.as_str())
}
}
impl Serialize for Format {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
serializer.serialize_str(self.as_str())
}
}
impl<'de> Deserialize<'de> for Format {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let name = String::deserialize(deserializer)?;
Self::named(&name).ok_or_else(|| {
de::Error::custom(format!(
"`{name}` is not a protocol family ({})",
Format::names()
))
})
}
}
#[derive(Debug, Clone, Copy)]
enum Registered {
OpenAi(openai::wire::Dialect),
Anthropic(anthropic::wire::Dialect),
Gemini,
}
impl PartialEq for Registered {
fn eq(&self, other: &Self) -> bool {
self.vendor() == other.vendor() && self.format() == other.format()
}
}
impl Registered {
fn vendor(&self) -> &'static str {
match self {
Self::OpenAi(dialect) => dialect.name,
Self::Anthropic(dialect) => dialect.name,
Self::Gemini => gemini::PROVIDER_NAME,
}
}
fn format(&self) -> Format {
match self {
Self::OpenAi(_) => Format::OpenAi,
Self::Anthropic(_) => Format::Anthropic,
Self::Gemini => Format::Gemini,
}
}
fn config(&self, api_key: impl Into<Secret>) -> ProviderConfig {
match self {
Self::OpenAi(dialect) => {
ProviderConfig::OpenAi(openai::wire::OpenAIConfig::with_key(dialect, api_key))
}
Self::Anthropic(dialect) => ProviderConfig::Anthropic(
anthropic::wire::AnthropicConfig::with_key(dialect, api_key),
),
Self::Gemini => ProviderConfig::Gemini(gemini::GeminiConfig::new(api_key)),
}
}
fn config_from_env(&self) -> Result<ProviderConfig, ConnectError> {
let tried = self.key_envs();
let (api_key, name) = key_from_env(self.vendor(), tried, self.requires_credential())?;
Ok(match self {
Self::OpenAi(dialect) => {
let auth = openai_auth(dialect, name);
ProviderConfig::OpenAi(openai::wire::OpenAIConfig::from_env_with_credential(
dialect, api_key, auth,
)?)
}
Self::Anthropic(dialect) => ProviderConfig::Anthropic(
anthropic::wire::AnthropicConfig::from_env_with_credential(dialect, api_key)?,
),
Self::Gemini => ProviderConfig::Gemini(gemini::GeminiConfig::new(api_key)),
})
}
fn key_envs(&self) -> Vec<&'static str> {
match self {
Self::OpenAi(dialect) => openai_key_envs(dialect),
Self::Anthropic(dialect) => vec![dialect.api_key_env],
Self::Gemini => vec![gemini::API_KEY_ENV],
}
}
fn requires_credential(&self) -> bool {
match self {
Self::OpenAi(dialect) => openai_requires_credential(dialect),
Self::Anthropic(_) | Self::Gemini => true,
}
}
}
fn openai_key_envs(dialect: &openai::wire::Dialect) -> Vec<&'static str> {
std::iter::once(dialect.api_key_env)
.chain(
dialect
.alternate_auth
.map(|alternative| alternative.api_key_env),
)
.collect()
}
fn openai_requires_credential(dialect: &openai::wire::Dialect) -> bool {
!matches!(dialect.quirks.auth, openai::wire::Auth::OptionalBearer)
}
fn openai_auth(dialect: &openai::wire::Dialect, name: Option<&'static str>) -> openai::wire::Auth {
match dialect.alternate_auth {
Some(alternative) if name == Some(alternative.api_key_env) => alternative.auth,
_ => dialect.quirks.auth,
}
}
fn key_from_env(
vendor: &str,
tried: Vec<&'static str>,
required: bool,
) -> Result<(String, Option<&'static str>), ConnectError> {
for name in &tried {
if let Some(value) = env::optional(name)?.filter(|value| !value.is_empty()) {
return Ok((value, Some(*name)));
}
}
if required {
return Err(ConnectError::MissingKey {
vendor: vendor.to_owned(),
tried: tried.into_iter().map(Cow::Borrowed).collect(),
});
}
Ok((String::new(), None))
}
#[derive(Debug)]
struct CatalogOnly {
vendor: &'static str,
home: &'static str,
credential: bool,
}
static CATALOG_ONLY: [CatalogOnly; 5] = [
CatalogOnly {
vendor: "aws_bedrock",
home: "the `rig-bedrock` crate",
credential: true,
},
CatalogOnly {
vendor: "vertexai",
home: "the `rig-vertexai` crate",
credential: true,
},
CatalogOnly {
vendor: "gemini-grpc",
home: "the `rig-gemini-grpc` crate",
credential: true,
},
CatalogOnly {
vendor: "candle",
home: "the `rig-candle` crate",
credential: false,
},
CatalogOnly {
vendor: "voyageai",
home: "`rig_core::providers::voyageai`, which serves embeddings and reranking only",
credential: true,
},
];
#[derive(Debug, Clone, Copy)]
enum Kind {
Registered(Registered),
CatalogOnly(&'static CatalogOnly),
Declared(&'static str),
}
static DECLARED: std::sync::RwLock<Vec<&'static str>> = std::sync::RwLock::new(Vec::new());
fn declared(vendor: &str) -> Option<&'static str> {
DECLARED
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.iter()
.copied()
.find(|name| *name == vendor)
}
#[derive(Debug, Clone, Copy)]
pub struct ProviderId(Kind);
impl PartialEq for ProviderId {
fn eq(&self, other: &Self) -> bool {
self.vendor() == other.vendor() && self.format() == other.format()
}
}
impl Eq for ProviderId {}
impl Hash for ProviderId {
fn hash<H: Hasher>(&self, state: &mut H) {
self.vendor().hash(state);
self.format().hash(state);
}
}
impl ProviderId {
pub fn all() -> impl Iterator<Item = ProviderId> {
openai::wire::all()
.map(|dialect| Registered::OpenAi(*dialect))
.chain(anthropic::wire::all().map(|dialect| Registered::Anthropic(*dialect)))
.chain(std::iter::once(Registered::Gemini))
.map(|registered| ProviderId(Kind::Registered(registered)))
}
pub fn new(vendor: &str, format: Format) -> Option<Self> {
let registered = match format {
Format::OpenAi => {
openai::wire::by_name(vendor).map(|dialect| Registered::OpenAi(*dialect))
}
Format::Anthropic => {
anthropic::wire::Dialect::by_name(vendor).map(Registered::Anthropic)
}
Format::Gemini => (vendor == gemini::PROVIDER_NAME).then_some(Registered::Gemini),
};
registered.map(|registered| Self(Kind::Registered(registered)))
}
pub fn catalog(vendor: &str) -> Option<Self> {
let primary = PRIMARY
.iter()
.find(|(name, _)| *name == vendor)
.and_then(|(_, format)| Self::new(vendor, *format));
primary
.or_else(|| Self::vendor_selections(vendor).next())
.or_else(|| {
CATALOG_ONLY
.iter()
.find(|provider| provider.vendor == vendor)
.map(|provider| Self(Kind::CatalogOnly(provider)))
})
.or_else(|| declared(vendor).map(|name| Self(Kind::Declared(name))))
}
pub fn declare(vendor: &str) -> Result<Self, SelectionError> {
if vendor.is_empty() || vendor.contains('/') || vendor.contains(char::is_whitespace) {
return Err(SelectionError::Malformed {
selection: vendor.to_owned(),
});
}
if let Some(known) = Self::catalog(vendor) {
return Ok(known);
}
let mut names = DECLARED
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let name = match names.iter().copied().find(|name| *name == vendor) {
Some(name) => name,
None => {
let name: &'static str = Box::leak(vendor.to_owned().into_boxed_str());
names.push(name);
name
}
};
Ok(Self(Kind::Declared(name)))
}
pub(crate) fn catalog_vendors() -> impl Iterator<Item = &'static str> {
let declared: Vec<&'static str> = DECLARED
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone();
Self::all()
.map(|id| id.vendor())
.chain(CATALOG_ONLY.iter().map(|provider| provider.vendor))
.chain(declared)
}
pub fn vendor(&self) -> &'static str {
match &self.0 {
Kind::Registered(registered) => registered.vendor(),
Kind::CatalogOnly(provider) => provider.vendor,
Kind::Declared(vendor) => vendor,
}
}
pub fn format(&self) -> Option<Format> {
match &self.0 {
Kind::Registered(registered) => Some(registered.format()),
Kind::CatalogOnly(_) | Kind::Declared(_) => None,
}
}
pub fn is_registered(&self) -> bool {
matches!(self.0, Kind::Registered(_))
}
pub fn is_declared(&self) -> bool {
matches!(self.0, Kind::Declared(_))
}
pub fn vendor_selections(vendor: &str) -> impl Iterator<Item = ProviderId> + '_ {
Self::all().filter(move |id| id.vendor() == vendor)
}
pub fn resolve(selection: &str) -> Result<Self, SelectionError> {
let malformed = || SelectionError::Malformed {
selection: selection.to_owned(),
};
let (vendor, format) = match selection.split_once('/') {
Some((_, rest)) if rest.contains('/') => return Err(malformed()),
Some((vendor, format)) => (vendor, Some(format)),
None => (selection, None),
};
if vendor.is_empty() {
return Err(malformed());
}
let Some(format) = format else {
let mut registered = Self::vendor_selections(vendor);
let first = registered.next().ok_or_else(|| SelectionError::Unknown {
vendor: vendor.to_owned(),
})?;
return match registered.next() {
None => Ok(first),
Some(_) => Err(SelectionError::Ambiguous {
vendor: vendor.to_owned(),
alternatives: alternatives(vendor),
}),
};
};
let Some(family) = Format::named(format) else {
return Err(SelectionError::UnknownFormat {
format: format.to_owned(),
});
};
Self::selection(vendor, family)
}
pub fn config(&self, api_key: impl Into<Secret>) -> Option<ProviderConfig> {
match &self.0 {
Kind::Registered(registered) => Some(registered.config(api_key)),
Kind::CatalogOnly(_) | Kind::Declared(_) => None,
}
}
pub fn api_key_envs(&self) -> Vec<&'static str> {
match &self.0 {
Kind::Registered(registered) => registered.key_envs(),
Kind::CatalogOnly(_) | Kind::Declared(_) => Vec::new(),
}
}
pub fn requires_credential(&self) -> bool {
match &self.0 {
Kind::Registered(registered) => registered.requires_credential(),
Kind::CatalogOnly(provider) => provider.credential,
Kind::Declared(_) => false,
}
}
fn registered(self) -> Result<Registered, ConnectError> {
match self.0 {
Kind::Registered(registered) => Ok(registered),
Kind::CatalogOnly(provider) => Err(ConnectError::CatalogOnly {
vendor: provider.vendor.to_owned(),
served_by: provider.home,
}),
Kind::Declared(vendor) => Err(ConnectError::CatalogOnly {
vendor: vendor.to_owned(),
served_by: "whoever declared it: provide it with `Connector::provide`",
}),
}
}
fn selection(vendor: &str, format: Format) -> Result<Self, SelectionError> {
Self::new(vendor, format).ok_or_else(|| {
let alternatives = alternatives(vendor);
match alternatives.is_empty() {
true => SelectionError::Unknown {
vendor: vendor.to_owned(),
},
false => SelectionError::Unregistered {
vendor: vendor.to_owned(),
format,
alternatives,
},
}
})
}
}
const PRIMARY: [(&str, Format); 4] = [
("zai", Format::OpenAi),
("minimax", Format::Anthropic),
("moonshot", Format::OpenAi),
("xiaomimimo", Format::OpenAi),
];
fn alternatives(vendor: &str) -> Vec<String> {
ProviderId::vendor_selections(vendor)
.map(|id| id.to_string())
.collect()
}
impl fmt::Display for ProviderId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self.format() {
Some(format) => write!(f, "{}/{format}", self.vendor()),
None => f.write_str(self.vendor()),
}
}
}
impl Serialize for ProviderId {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
serializer.collect_str(self)
}
}
impl<'de> Deserialize<'de> for ProviderId {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let selection = String::deserialize(deserializer)?;
Self::resolve(&selection)
.or_else(|error| {
Self::catalog(&selection)
.filter(|id| !id.is_registered())
.ok_or(error)
})
.map_err(de::Error::custom)
}
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum SelectionError {
#[error("no registered provider is named `{vendor}`")]
Unknown {
vendor: String,
},
#[error(
"`{vendor}` speaks no {format} endpoint in this build (it speaks {})",
alternatives.join(", ")
)]
Unregistered {
vendor: String,
format: Format,
alternatives: Vec<String>,
},
#[error(
"`{vendor}` names more than one registered selection: name one of {}",
alternatives.join(", ")
)]
Ambiguous {
vendor: String,
alternatives: Vec<String>,
},
#[error("`{format}` is not a protocol family ({})", Format::names())]
UnknownFormat {
format: String,
},
#[error("`{selection}` is not a provider selection: expected `vendor` or `vendor/format`")]
Malformed {
selection: String,
},
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum ProviderConfig {
#[serde(rename = "openai")]
OpenAi(openai::wire::OpenAIConfig),
#[serde(rename = "anthropic")]
Anthropic(anthropic::wire::AnthropicConfig),
#[serde(rename = "gemini")]
Gemini(gemini::GeminiConfig),
}
impl ProviderConfig {
pub fn id(&self) -> Option<ProviderId> {
ProviderId::new(self.vendor(), self.format())
}
pub fn vendor(&self) -> &'static str {
match self {
Self::OpenAi(provider) => provider.dialect.name,
Self::Anthropic(provider) => provider.dialect.name,
Self::Gemini(_) => gemini::PROVIDER_NAME,
}
}
pub fn format(&self) -> Format {
match self {
Self::OpenAi(_) => Format::OpenAi,
Self::Anthropic(_) => Format::Anthropic,
Self::Gemini(_) => Format::Gemini,
}
}
pub fn is_unauthenticated(&self) -> bool {
match self {
Self::OpenAi(provider) => provider.api_key.is_empty(),
Self::Anthropic(provider) => provider.api_key.is_empty(),
Self::Gemini(provider) => provider.api_key.is_empty(),
}
}
pub fn with_credential(mut self, api_key: impl Into<Secret>) -> Self {
let api_key = api_key.into();
match &mut self {
Self::OpenAi(provider) => provider.api_key = api_key,
Self::Anthropic(provider) => provider.api_key = api_key,
Self::Gemini(provider) => provider.api_key = api_key,
}
self
}
pub fn base_url(&self) -> &str {
match self {
Self::OpenAi(provider) => &provider.base_url,
Self::Anthropic(provider) => &provider.base_url,
Self::Gemini(provider) => &provider.base_url,
}
}
pub fn with_base_url(self, base_url: impl Into<String>) -> Self {
let base_url = base_url.into();
match self {
Self::OpenAi(provider) => Self::OpenAi(provider.with_base_url(base_url)),
Self::Anthropic(provider) => Self::Anthropic(provider.with_base_url(base_url)),
Self::Gemini(provider) => Self::Gemini(provider.with_base_url(base_url)),
}
}
pub fn completion_handler(
&self,
label: &str,
model: &str,
http: DynHttpClient,
facts: ModelFacts,
) -> ErasedHandler {
ErasedHandler::new(ModelAdapter::new(
label,
self.completion_model(model, http, facts),
))
}
pub fn completion_model(
&self,
model: &str,
http: DynHttpClient,
facts: ModelFacts,
) -> DynModel<Completion> {
match self {
Self::OpenAi(provider) => {
let mut model = provider.clone().connect(http).completion(model);
model.wire = model.wire.with_facts(facts);
model.erase()
}
Self::Anthropic(provider) => {
let mut model = provider.clone().connect(http).completion(model);
model.wire = model.wire.with_facts(facts);
model.erase()
}
Self::Gemini(provider) => {
let mut model = provider.clone().connect(http).completion(model);
model.wire = model.wire.with_facts(facts);
model.erase()
}
}
}
}
#[non_exhaustive]
#[derive(Clone, Copy, Debug)]
pub enum ModelSelector<'a> {
Spec(&'a ModelSpec),
Reference(&'a str),
}
impl<'a> From<&'a ModelSpec> for ModelSelector<'a> {
fn from(spec: &'a ModelSpec) -> Self {
Self::Spec(spec)
}
}
impl<'a> From<&'a str> for ModelSelector<'a> {
fn from(reference: &'a str) -> Self {
Self::Reference(reference)
}
}
impl<'a> From<&'a String> for ModelSelector<'a> {
fn from(reference: &'a String) -> Self {
Self::Reference(reference)
}
}
#[non_exhaustive]
#[derive(Clone, Default)]
pub struct ConnectOptions {
api_key: Option<Secret>,
account_id: Option<String>,
pub(crate) http: Option<DynHttpClient>,
format: Option<Format>,
base_url: Option<String>,
}
impl ConnectOptions {
pub fn new() -> Self {
Self::default()
}
pub fn api_key(mut self, api_key: impl Into<Secret>) -> Self {
self.api_key = Some(api_key.into());
self
}
pub fn account_id(mut self, account_id: impl Into<String>) -> Self {
self.account_id = Some(account_id.into());
self
}
pub fn http(mut self, http: impl HttpClientExt + 'static) -> Self {
self.http = Some(DynHttpClient::new(http));
self
}
pub fn format(mut self, format: Format) -> Self {
self.format = Some(format);
self
}
pub fn base_url(mut self, base_url: impl Into<String>) -> Self {
self.base_url = Some(base_url.into());
self
}
}
impl fmt::Debug for ConnectOptions {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ConnectOptions")
.field("api_key", &self.api_key.as_ref().map(|_| "<set>"))
.field("account_id", &self.account_id.as_ref().map(|_| "<set>"))
.field("http", &self.http.as_ref().map(|_| "<set>"))
.field("format", &self.format)
.field("base_url", &self.base_url)
.finish()
}
}
#[non_exhaustive]
#[derive(Debug, thiserror::Error)]
pub enum ConnectError {
#[error("`{reference}` names no model: expected `vendor/model`")]
Malformed {
reference: String,
},
#[error("the registry cannot connect to `{vendor}`: its models are served by {served_by}")]
CatalogOnly {
vendor: String,
served_by: &'static str,
},
#[error(transparent)]
NotFound(#[from] NotFound),
#[error(transparent)]
Selection(#[from] SelectionError),
#[error(
"no key for `{vendor}`: set {} or pass one with `ConnectOptions::api_key`",
tried.iter().map(|name| format!("`{name}`")).collect::<Vec<_>>().join(" or ")
)]
MissingKey {
vendor: String,
tried: Vec<Cow<'static, str>>,
},
#[error("{0}")]
Provider(String),
#[error(transparent)]
Env(#[from] EnvError),
#[error("no HTTP client: pass one with `ConnectOptions::http`")]
NoHttpClient,
}
impl Catalog {
#[cfg(feature = "reqwest")]
#[cfg_attr(docsrs, doc(cfg(feature = "reqwest")))]
pub fn connect<'a>(
&self,
model: impl Into<ModelSelector<'a>>,
) -> Result<DynModel<Completion>, ConnectError> {
self.connect_with(model, ConnectOptions::new())
}
pub fn connect_with<'a>(
&self,
model: impl Into<ModelSelector<'a>>,
options: ConnectOptions,
) -> Result<DynModel<Completion>, ConnectError> {
let (provider, model, facts) = match model.into() {
ModelSelector::Spec(spec) => (
spec.provider,
spec.id.as_str(),
ModelFacts::new(spec.clone()).with_catalog(self),
),
ModelSelector::Reference(text) => {
let reference = Reference::parse(text).ok_or_else(|| ConnectError::Malformed {
reference: text.to_owned(),
})?;
let vendor =
ProviderId::catalog(reference.vendor).ok_or_else(|| self.not_found(text))?;
let provider = self
.get_vendor(reference.vendor, reference.model)
.map_or(vendor, |resolved| resolved.spec.provider);
let provider = match reference.format {
Some(format) if provider.is_registered() => {
ProviderId::selection(reference.vendor, format)?
}
_ => provider,
};
(
provider,
reference.model,
self.facts(provider, reference.model),
)
}
};
let provider = match options.format {
Some(format) if provider.is_registered() => {
ProviderId::selection(provider.vendor(), format)?
}
_ => provider,
};
let registered = provider.registered()?;
let config = match options.api_key {
Some(api_key) => registered.config(api_key),
None => registered.config_from_env()?,
};
let config = match (config, options.account_id) {
(ProviderConfig::OpenAi(provider), Some(account_id)) => {
ProviderConfig::OpenAi(provider.with_account_id(account_id))
}
(config, _) => config,
};
let config = match options.base_url {
Some(base_url) => config.with_base_url(base_url),
None => config,
};
let http = match options.http {
Some(http) => http,
None => shared_http()?,
};
Ok(config.completion_model(model, http, facts))
}
}
#[cfg(feature = "reqwest")]
fn shared_http() -> Result<DynHttpClient, ConnectError> {
Ok(rig_reqwest::shared())
}
#[cfg(not(feature = "reqwest"))]
fn shared_http() -> Result<DynHttpClient, ConnectError> {
Err(ConnectError::NoHttpClient)
}
#[cfg(test)]
mod tests;