diff options
Diffstat (limited to 'src/client/model.rs')
| -rw-r--r-- | src/client/model.rs | 124 |
1 files changed, 83 insertions, 41 deletions
diff --git a/src/client/model.rs b/src/client/model.rs index 5d496d2..47f6598 100644 --- a/src/client/model.rs +++ b/src/client/model.rs @@ -1,7 +1,7 @@ use super::{ - list_chat_models, list_embedding_models, list_reranker_models, + list_all_models, list_client_names, message::{Message, MessageContent, MessageContentPart}, - MessageContentToolCalls, + ApiPatch, MessageContentToolCalls, RequestPatch, }; use crate::config::Config; @@ -9,6 +9,7 @@ use crate::utils::{estimate_token_length, format_option_value}; use anyhow::{bail, Result}; use serde::{Deserialize, Serialize}; +use std::fmt::Display; const PER_MESSAGES_TOKENS: usize = 5; const BASIS_TOKENS: usize = 2; @@ -43,29 +44,8 @@ impl Model { .collect() } - pub fn retrieve_chat(config: &Config, model_id: &str) -> Result<Self> { - match Self::find(&list_chat_models(config), model_id) { - Some(v) => Ok(v), - None => bail!("Unknown chat model '{model_id}'"), - } - } - - pub fn retrieve_embedding(config: &Config, model_id: &str) -> Result<Self> { - match Self::find(&list_embedding_models(config), model_id) { - Some(v) => Ok(v), - None => bail!("Unknown embedding model '{model_id}'"), - } - } - - pub fn retrieve_reranker(config: &Config, model_id: &str) -> Result<Self> { - match Self::find(&list_reranker_models(config), model_id) { - Some(v) => Ok(v), - None => bail!("Unknown reranker model '{model_id}'"), - } - } - - pub fn find(models: &[&Self], model_id: &str) -> Option<Self> { - let mut model = None; + pub fn retrieve_model(config: &Config, model_id: &str, model_type: ModelType) -> Result<Self> { + let models = list_all_models(config); let (client_name, model_name) = match model_id.split_once(':') { Some((client_name, model_name)) => { if model_name.is_empty() { @@ -78,21 +58,33 @@ impl Model { }; match model_name { Some(model_name) => { - if let Some(found) = models.iter().find(|v| v.id() == model_id) { - model = Some((*found).clone()); - } else if let Some(found) = models.iter().find(|v| v.client_name == client_name) { - let mut found = (*found).clone(); - found.data.name = model_name.to_string(); - model = Some(found) + if let Some(model) = models.iter().find(|v| v.id() == model_id) { + if model.model_type() == model_type { + return Ok((*model).clone()); + } else { + bail!("Model '{model_id}' is not a {model_type} model") + } + } + if list_client_names(config) + .into_iter() + .any(|v| *v == client_name) + && model_type.can_create_from_name() + { + let mut new_model = Self::new(client_name, model_name); + new_model.data.model_type = model_type.to_string(); + return Ok(new_model); } } None => { - if let Some(found) = models.iter().find(|v| v.client_name == client_name) { - model = Some((*found).clone()); + if let Some(found) = models + .iter() + .find(|v| v.client_name == client_name && v.model_type() == model_type) + { + return Ok((*found).clone()); } } - } - model + }; + bail!("Unknown {model_type} model '{model_id}'") } pub fn id(&self) -> String { @@ -111,8 +103,14 @@ impl Model { &self.data.name } - pub fn model_type(&self) -> &str { - &self.data.model_type + pub fn model_type(&self) -> ModelType { + if self.data.model_type.starts_with("embed") { + ModelType::Embedding + } else if self.data.model_type.starts_with("rerank") { + ModelType::Reranker + } else { + ModelType::Chat + } } pub fn data(&self) -> &ModelData { @@ -125,7 +123,7 @@ impl Model { pub fn description(&self) -> String { match self.model_type() { - "chat" => { + ModelType::Chat => { let ModelData { max_input_tokens, max_output_tokens, @@ -156,7 +154,7 @@ impl Model { max_input_tokens, max_output_tokens, input_price, output_price, capabilities ) } - "embedding" => { + ModelType::Embedding => { let ModelData { input_price, max_tokens_per_chunk, @@ -168,7 +166,7 @@ impl Model { let price = format_option_value(input_price); format!("max-tokens:{max_tokens};max-batch:{max_batch};price:{price}") } - _ => String::new(), + ModelType::Reranker => String::new(), } } @@ -310,13 +308,14 @@ impl ModelData { pub fn new(name: &str) -> Self { Self { name: name.to_string(), + model_type: default_model_type(), ..Default::default() } } } #[derive(Debug, Clone, Deserialize)] -pub struct BuiltinModels { +pub struct PredefinedModels { pub platform: String, pub models: Vec<ModelData>, } @@ -324,3 +323,46 @@ pub struct BuiltinModels { fn default_model_type() -> String { "chat".into() } + +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum ModelType { + Chat, + Embedding, + Reranker, +} + +impl Display for ModelType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ModelType::Chat => write!(f, "chat"), + ModelType::Embedding => write!(f, "embedding"), + ModelType::Reranker => write!(f, "reranker"), + } + } +} + +impl ModelType { + pub fn can_create_from_name(self) -> bool { + match self { + ModelType::Chat => true, + ModelType::Embedding => false, + ModelType::Reranker => true, + } + } + + pub fn api_name(self) -> &'static str { + match self { + ModelType::Chat => "chat_completions", + ModelType::Embedding => "embeddings", + ModelType::Reranker => "rerank", + } + } + + pub fn extract_patch(self, patch: &RequestPatch) -> Option<&ApiPatch> { + match self { + ModelType::Chat => patch.chat_completions.as_ref(), + ModelType::Embedding => patch.embeddings.as_ref(), + ModelType::Reranker => patch.rerank.as_ref(), + } + } +} |