refactor: set tokens count factors (#195)

1 files changed, 11 insertions, 6 deletions

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<String>,
@@ -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)
-}