From 5c7bfd92ff3e557477969be9db0638ef0d3d3659 Mon Sep 17 00:00:00 2001 From: sigoden Date: Thu, 2 Nov 2023 07:08:54 +0800 Subject: refactor: set tokens count factors (#195) --- src/client/openai.rs | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) (limited to 'src/client/openai.rs') diff --git a/src/client/openai.rs b/src/client/openai.rs index 387baf0..a82ee9a 100644 --- a/src/client/openai.rs +++ b/src/client/openai.rs @@ -1,4 +1,7 @@ -use super::{ExtraConfig, ModelInfo, OpenAIClient, PromptKind, PromptType, SendData, ReplyStreamHandler}; +use super::{ + ExtraConfig, ModelInfo, OpenAIClient, PromptKind, PromptType, ReplyStreamHandler, SendData, + TokensCountFactors, +}; use anyhow::{anyhow, bail, Result}; use async_trait::async_trait; @@ -18,6 +21,8 @@ const MODELS: [(&str, usize); 4] = [ ("gpt-4-32k", 32768), ]; +pub const OPENAI_TOKENS_COUNT_FACTORS: TokensCountFactors = (5, 2); + #[derive(Debug, Clone, Deserialize, Default)] pub struct OpenAIConfig { pub name: Option, @@ -38,7 +43,11 @@ impl OpenAIClient { let client = Self::name(local_config); MODELS .into_iter() - .map(|(name, max_tokens)| openai_tokens_formula(ModelInfo::new(index, client, name).set_max_tokens(Some(max_tokens)))) + .map(|(name, max_tokens)| { + ModelInfo::new(index, client, name) + .set_max_tokens(Some(max_tokens)) + .set_tokens_count_factors(OPENAI_TOKENS_COUNT_FACTORS) + }) .collect() } @@ -135,7 +144,3 @@ pub fn openai_build_body(data: SendData, model: String) -> Value { } body } - -pub fn openai_tokens_formula(model: ModelInfo) -> ModelInfo { - model.set_tokens_formula(5, 2) -} -- cgit v1.2.3