#[cfg(test)]
mod tests;
use std::{path::PathBuf, sync::Arc};
use mentra::{
BuiltinProvider, ProviderId,
provider_core::{
AuthScheme, Provider, StripModelPrefix, chat_completions,
gateway_ring::{GatewayRing, GatewayRingEvent, GatewayRingPolicy},
},
};
use crate::{error::RunError, provider, runtime::credential::Credential};
use super::{RuntimeBuilder, Wire, provider_settlement::responses_provider};
#[derive(Clone, PartialEq, Eq)]
pub struct GatewayMember {
base_url: String,
key: MemberKey,
strip_model_prefix: Option<String>,
}
#[derive(Clone, PartialEq, Eq)]
enum MemberKey {
None,
Value(String),
Env(String),
File(PathBuf),
}
impl GatewayMember {
pub fn new(base_url: impl Into<String>) -> Self {
Self {
base_url: base_url.into(),
key: MemberKey::None,
strip_model_prefix: None,
}
}
#[must_use]
pub fn with_api_key(self, api_key: impl Into<String>) -> Self {
Self {
key: MemberKey::Value(api_key.into()),
..self
}
}
#[must_use]
pub fn with_api_key_env(self, var: impl Into<String>) -> Self {
Self {
key: MemberKey::Env(var.into()),
..self
}
}
#[must_use]
pub fn with_api_key_file(self, path: impl Into<PathBuf>) -> Self {
Self {
key: MemberKey::File(path.into()),
..self
}
}
#[must_use]
pub fn with_model_prefix_stripped(self, prefix: impl Into<String>) -> Self {
Self {
strip_model_prefix: Some(prefix.into()),
..self
}
}
pub fn base_url(&self) -> &str {
&self.base_url
}
fn resolve_key(
&self,
index: usize,
lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<Option<String>, provider::ProviderError> {
let missing = |origin: String| provider::ProviderError::GatewayCredentialMissing {
member: index,
origin,
};
match &self.key {
MemberKey::None => Ok(None),
MemberKey::Value(key) => Ok(Some(key.clone())),
MemberKey::Env(var) => lookup(var)
.map(|key| key.trim().to_string())
.filter(|key| !key.is_empty())
.map(Some)
.ok_or_else(|| missing(format!("${var}"))),
MemberKey::File(path) => {
let key = std::fs::read_to_string(path).map_err(|source| {
provider::ProviderError::GatewayCredentialUnreadable {
member: index,
path: path.clone(),
source,
}
})?;
let key = key.trim();
if key.is_empty() {
return Err(missing(path.display().to_string()));
}
Ok(Some(key.to_string()))
}
}
}
}
impl std::fmt::Debug for GatewayMember {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let key = match &self.key {
MemberKey::None => "none".to_string(),
MemberKey::Value(_) => "<redacted>".to_string(),
MemberKey::Env(var) => format!("${var}"),
MemberKey::File(path) => path.display().to_string(),
};
f.debug_struct("GatewayMember")
.field("base_url", &self.base_url)
.field("api_key", &key)
.field("strip_model_prefix", &self.strip_model_prefix)
.finish()
}
}
type Observer = Arc<dyn Fn(GatewayRingEvent) + Send + Sync>;
#[derive(Clone, Default)]
pub(in crate::runtime) struct GatewayRingSpec {
pub(super) members: Option<Vec<GatewayMember>>,
pub(super) policy: Option<GatewayRingPolicy>,
pub(super) observer: Option<Observer>,
}
impl GatewayRingSpec {
pub(in crate::runtime) fn is_stated(&self) -> bool {
self.members.is_some()
}
pub(in crate::runtime) fn filled_from(self, ring: &crate::GatewayRingConfig) -> Self {
Self {
members: Some(ring.members.clone()),
policy: self.policy.or(ring.policy),
observer: self.observer,
}
}
}
impl std::fmt::Debug for GatewayRingSpec {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("GatewayRingSpec")
.field("members", &self.members)
.field("policy", &self.policy)
.field("observer", &self.observer.as_ref().map(|_| "<observer>"))
.finish()
}
}
impl RuntimeBuilder {
#[must_use]
pub fn with_gateway_ring(self, members: impl IntoIterator<Item = GatewayMember>) -> Self {
let spec = GatewayRingSpec {
members: Some(members.into_iter().collect()),
..self.gateway_ring.clone().unwrap_or_default()
};
Self {
gateway_ring: Some(spec),
..self
}
}
#[must_use]
pub fn with_gateway_ring_policy(self, policy: GatewayRingPolicy) -> Self {
let spec = GatewayRingSpec {
policy: Some(policy),
..self.gateway_ring.clone().unwrap_or_default()
};
Self {
gateway_ring: Some(spec),
..self
}
}
#[must_use]
pub fn with_gateway_ring_observer(
self,
observer: impl Fn(GatewayRingEvent) + Send + Sync + 'static,
) -> Self {
let spec = GatewayRingSpec {
observer: Some(Arc::new(observer)),
..self.gateway_ring.clone().unwrap_or_default()
};
Self {
gateway_ring: Some(spec),
..self
}
}
}
pub(super) fn assemble(
spec: GatewayRingSpec,
provider: BuiltinProvider,
wire: Wire,
) -> Result<GatewayRing, RunError> {
let lookup = |var: &str| std::env::var(var).ok();
let members = spec
.members
.unwrap_or_default()
.iter()
.enumerate()
.map(|(index, member)| build_member(index, member, provider, wire, &lookup))
.collect::<Result<Vec<_>, _>>()?;
let ring = GatewayRing::new(members, spec.policy.unwrap_or_default())
.map_err(|_| provider::ProviderError::EmptyGatewayRing)?;
Ok(match spec.observer {
Some(observer) => ring.with_observer(move |event| observer(event)),
None => ring,
})
}
fn build_member(
index: usize,
member: &GatewayMember,
provider: BuiltinProvider,
wire: Wire,
lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<Arc<dyn Provider>, RunError> {
let base_url = provider::normalize_base_url(&member.base_url)?;
let api_key = member.resolve_key(index, lookup)?;
let credential = Credential::new(api_key.as_deref());
let built: Arc<dyn Provider> = match wire {
Wire::Responses => Arc::new(responses_provider(provider, &base_url, credential)),
Wire::ChatCompletions => {
Arc::new(chat_completions_provider(provider, &base_url, credential))
}
};
Ok(match &member.strip_model_prefix {
Some(prefix) => Arc::new(StripModelPrefix::new(built, prefix.as_str())),
None => built,
})
}
fn chat_completions_provider(
provider: BuiltinProvider,
base_url: &str,
credential: Credential,
) -> chat_completions::ChatCompletionsProvider<Credential> {
let mut definition = chat_completions::definition(ProviderId::from(provider), base_url);
definition.descriptor.display_name = Some(format!("OpenAI-compatible ({base_url})"));
if !credential.is_some() {
definition.auth_scheme = AuthScheme::None;
}
chat_completions::ChatCompletionsProvider::new(definition, credential)
}