use crate::core::CapabilitySupport;
use crate::engine::ModelCapabilities;
#[cfg(test)]
#[path = "../tests/client/endpoint_tests.rs"]
mod endpoint_tests;
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum EndpointRef {
Local {
id: String,
},
Gateway {
id: String,
},
Provider {
id: String,
},
}
impl EndpointRef {
pub fn id(&self) -> &str {
match self {
Self::Local { id } | Self::Gateway { id } | Self::Provider { id } => id,
}
}
pub const fn kind(&self) -> &'static str {
match self {
Self::Local { .. } => "local",
Self::Gateway { .. } => "gateway",
Self::Provider { .. } => "provider",
}
}
pub(crate) fn is_local(&self) -> bool {
matches!(self, Self::Local { .. })
}
pub fn gateway(id: impl Into<String>) -> Self {
Self::Gateway { id: id.into() }
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct EndpointCapabilities {
pub query: CapabilitySupport,
pub chat: CapabilitySupport,
pub embed: CapabilitySupport,
}
impl EndpointCapabilities {
pub(crate) fn from_local(capabilities: &ModelCapabilities) -> Self {
Self {
query: support(capabilities.supports_text_generation),
chat: support(capabilities.supports_text_generation && capabilities.has_chat_template),
embed: support(capabilities.supports_embeddings),
}
}
pub const fn unknown() -> Self {
Self {
query: CapabilitySupport::Unknown,
chat: CapabilitySupport::Unknown,
embed: CapabilitySupport::Unknown,
}
}
pub fn for_operation(&self, operation: &'static str) -> CapabilitySupport {
match operation {
"query" => self.query,
"chat" => self.chat,
"embed" => self.embed,
_ => CapabilitySupport::Unsupported,
}
}
}
const fn support(enabled: bool) -> CapabilitySupport {
if enabled {
CapabilitySupport::Supported
} else {
CapabilitySupport::Unsupported
}
}