refactor: set tokens count factors (#195)

4 files changed, 28 insertions, 11 deletions

diff --git a/src/client/azure_openai.rs b/src/client/azure_openai.rs
index 3a48145..5603666 100644
--- a/src/client/azure_openai.rs
+++ b/src/client/azure_openai.rs
@@ -1,4 +1,4 @@
-use super::openai::{openai_build_body, openai_tokens_formula};
+use super::openai::{openai_build_body, OPENAI_TOKENS_COUNT_FACTORS};
 use super::{AzureOpenAIClient, ExtraConfig, ModelInfo, PromptKind, PromptType, SendData};
 
 use anyhow::{anyhow, Result};
@@ -46,7 +46,11 @@ impl AzureOpenAIClient {
         local_config
             .models
             .iter()
-            .map(|v| openai_tokens_formula(ModelInfo::new(index, client, &v.name).set_max_tokens(v.max_tokens)))
+            .map(|v| {
+                ModelInfo::new(index, client, &v.name)
+                    .set_max_tokens(v.max_tokens)
+                    .set_tokens_count_factors(OPENAI_TOKENS_COUNT_FACTORS)
+            })
             .collect()
     }
 
diff --git a/src/client/localai.rs b/src/client/localai.rs
index 026ae27..d438388 100644
--- a/src/client/localai.rs
+++ b/src/client/localai.rs
@@ -1,4 +1,4 @@
-use super::openai::{openai_build_body, openai_tokens_formula};
+use super::openai::{openai_build_body, OPENAI_TOKENS_COUNT_FACTORS};
 use super::{ExtraConfig, LocalAIClient, ModelInfo, PromptKind, PromptType, SendData};
 
 use anyhow::Result;
@@ -45,7 +45,11 @@ impl LocalAIClient {
         local_config
             .models
             .iter()
-            .map(|v| openai_tokens_formula(ModelInfo::new(index, client, &v.name).set_max_tokens(v.max_tokens)))
+            .map(|v| {
+                ModelInfo::new(index, client, &v.name)
+                    .set_max_tokens(v.max_tokens)
+                    .set_tokens_count_factors(OPENAI_TOKENS_COUNT_FACTORS)
+            })
             .collect()
     }
 
diff --git a/src/client/mod.rs b/src/client/mod.rs
index dba049d..5fa6146 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -3,7 +3,11 @@ mod common;
 
 pub use common::*;
 
-use crate::{config::ModelInfo, repl::ReplyStreamHandler, utils::PromptKind};
+use crate::{
+    config::{ModelInfo, TokensCountFactors},
+    repl::ReplyStreamHandler,
+    utils::PromptKind,
+};
 
 register_client!(
     (openai, "openai", OpenAI, OpenAIConfig, OpenAIClient),
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)
-}