summaryrefslogtreecommitdiffstats
path: root/src/client/macros.rs
diff options
context:
space:
mode:
Diffstat (limited to 'src/client/macros.rs')
-rw-r--r--src/client/macros.rs45
1 files changed, 27 insertions, 18 deletions
diff --git a/src/client/macros.rs b/src/client/macros.rs
index 0543cd6..4f52044 100644
--- a/src/client/macros.rs
+++ b/src/client/macros.rs
@@ -52,9 +52,10 @@ macro_rules! register_client {
pub fn list_models(local_config: &$config) -> Vec<Model> {
let client_name = Self::name(local_config);
if local_config.models.is_empty() {
- if let Some(models) = $crate::client::ALL_MODELS.iter().find(|v| {
+ if let Some(models) = $crate::client::ALL_PREDEFINED_MODELS.iter().find(|v| {
v.platform == $name ||
- ($name == OpenAICompatibleClient::NAME && local_config.name.as_deref() == Some(&v.platform))
+ ($name == OpenAICompatibleClient::NAME
+ && local_config.name.as_ref().map(|name| name.starts_with(&v.platform)).unwrap_or_default())
}) {
return Model::from_config(client_name, &models.models);
}
@@ -98,32 +99,40 @@ macro_rules! register_client {
anyhow::bail!("Unknown client '{}'", client)
}
- static ALL_CLIENT_MODELS: std::sync::OnceLock<Vec<$crate::client::Model>> = std::sync::OnceLock::new();
+ static ALL_CLIENT_NAMES: std::sync::OnceLock<Vec<String>> = std::sync::OnceLock::new();
- pub fn list_models(config: &$crate::config::Config) -> Vec<&'static $crate::client::Model> {
- let models = ALL_CLIENT_MODELS.get_or_init(|| {
+ pub fn list_client_names(config: &$crate::config::Config) -> Vec<&'static String> {
+ let names = ALL_CLIENT_NAMES.get_or_init(|| {
config
.clients
.iter()
.flat_map(|v| match v {
- $(ClientConfig::$config(c) => $client::list_models(c),)+
+ $(ClientConfig::$config(c) => vec![$client::name(c).to_string()],)+
ClientConfig::Unknown => vec![],
})
.collect()
});
- models.iter().collect()
+ names.iter().collect()
}
- pub fn list_chat_models(config: &$crate::config::Config) -> Vec<&'static $crate::client::Model> {
- list_models(config).into_iter().filter(|v| v.model_type() == "chat").collect()
- }
+ static ALL_MODELS: std::sync::OnceLock<Vec<$crate::client::Model>> = std::sync::OnceLock::new();
- pub fn list_embedding_models(config: &$crate::config::Config) -> Vec<&'static $crate::client::Model> {
- list_models(config).into_iter().filter(|v| v.model_type() == "embedding").collect()
+ pub fn list_all_models(config: &$crate::config::Config) -> Vec<&'static $crate::client::Model> {
+ let models = ALL_MODELS.get_or_init(|| {
+ config
+ .clients
+ .iter()
+ .flat_map(|v| match v {
+ $(ClientConfig::$config(c) => $client::list_models(c),)+
+ ClientConfig::Unknown => vec![],
+ })
+ .collect()
+ });
+ models.iter().collect()
}
- pub fn list_reranker_models(config: &$crate::config::Config) -> Vec<&'static $crate::client::Model> {
- list_models(config).into_iter().filter(|v| v.model_type() == "reranker").collect()
+ pub fn list_models(config: &$crate::config::Config, model_type: $crate::client::ModelType) -> Vec<&'static $crate::client::Model> {
+ list_all_models(config).into_iter().filter(|v| v.model_type() == model_type).collect()
}
};
}
@@ -175,7 +184,7 @@ macro_rules! impl_client_trait {
data: $crate::client::ChatCompletionsData,
) -> anyhow::Result<$crate::client::ChatCompletionsOutput> {
let request_data = $prepare_chat_completions(self, data)?;
- let builder = self.request_builder(client, request_data, ApiType::ChatCompletions);
+ let builder = self.request_builder(client, request_data);
$chat_completions(builder, self.model()).await
}
@@ -186,7 +195,7 @@ macro_rules! impl_client_trait {
data: $crate::client::ChatCompletionsData,
) -> Result<()> {
let request_data = $prepare_chat_completions(self, data)?;
- let builder = self.request_builder(client, request_data, ApiType::ChatCompletions);
+ let builder = self.request_builder(client, request_data);
$chat_completions_streaming(builder, handler, self.model()).await
}
@@ -196,7 +205,7 @@ macro_rules! impl_client_trait {
data: &$crate::client::EmbeddingsData,
) -> Result<$crate::client::EmbeddingsOutput> {
let request_data = $prepare_embeddings(self, data)?;
- let builder = self.request_builder(client, request_data, ApiType::Embeddings);
+ let builder = self.request_builder(client, request_data);
$embeddings(builder, self.model()).await
}
@@ -206,7 +215,7 @@ macro_rules! impl_client_trait {
data: &$crate::client::RerankData,
) -> Result<$crate::client::RerankOutput> {
let request_data = $prepare_rerank(self, data)?;
- let builder = self.request_builder(client, request_data, ApiType::Rerank);
+ let builder = self.request_builder(client, request_data);
$rerank(builder, self.model()).await
}
}