From fe35cfd9419302f01baf9672493c0b0a4b41d889 Mon Sep 17 00:00:00 2001 From: sigoden Date: Sat, 13 Jan 2024 19:52:07 +0800 Subject: feat: supports model capabilities (#297) 1. automatically switch to the model that has the necessary capabilities. 2. throw an error if the client does not have a model with the necessary capabilities --- src/client/gemini.rs | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) (limited to 'src/client/gemini.rs') diff --git a/src/client/gemini.rs b/src/client/gemini.rs index 98fd60d..6ff3f95 100644 --- a/src/client/gemini.rs +++ b/src/client/gemini.rs @@ -3,7 +3,7 @@ use super::{ SendData, TokensCountFactors, }; -use crate::{config::GlobalConfig, render::ReplyHandler, utils::PromptKind}; +use crate::{render::ReplyHandler, utils::PromptKind}; use anyhow::{anyhow, bail, Result}; use async_trait::async_trait; @@ -14,10 +14,10 @@ use serde_json::{json, Value}; const API_BASE: &str = "https://generativelanguage.googleapis.com/v1beta/models/"; -const MODELS: [(&str, usize); 3] = [ - ("gemini-pro", 32768), - ("gemini-pro-vision", 16384), - ("gemini-ultra", 32768), +const MODELS: [(&str, usize, &str); 3] = [ + ("gemini-pro", 32768, "text"), + ("gemini-pro-vision", 16384, "vision"), + ("gemini-ultra", 32768, "text"), ]; const TOKENS_COUNT_FACTORS: TokensCountFactors = (5, 2); @@ -31,9 +31,7 @@ pub struct GeminiConfig { #[async_trait] impl Client for GeminiClient { - fn config(&self) -> (&GlobalConfig, &Option) { - (&self.global_config, &self.config.extra) - } + client_common_fns!(); async fn send_message_inner(&self, client: &ReqwestClient, data: SendData) -> Result { let builder = self.request_builder(client, data)?; @@ -61,8 +59,9 @@ impl GeminiClient { let client_name = Self::name(local_config); MODELS .into_iter() - .map(|(name, max_tokens)| { + .map(|(name, max_tokens, capabilities)| { Model::new(client_name, name) + .set_capabilities(capabilities.into()) .set_max_tokens(Some(max_tokens)) .set_tokens_count_factors(TOKENS_COUNT_FACTORS) }) -- cgit v1.2.3