3 files changed, 9 insertions, 5 deletions
diff --git a/src/client/azure_openai.rs b/src/client/azure_openai.rs
index d0fa31f..3a48145 100644
--- a/src/client/azure_openai.rs
+++ b/src/client/azure_openai.rs
@@ -1,4 +1,4 @@
-use super::openai::openai_build_body;
+use super::openai::{openai_build_body, openai_tokens_formula};
 use super::{AzureOpenAIClient, ExtraConfig, ModelInfo, PromptKind, PromptType, SendData};
 
 use anyhow::{anyhow, Result};
@@ -46,7 +46,7 @@ impl AzureOpenAIClient {
         local_config
             .models
             .iter()
-            .map(|v| ModelInfo::new(client, &v.name, v.max_tokens, index))
+            .map(|v| openai_tokens_formula(ModelInfo::new(index, client, &v.name).set_max_tokens(v.max_tokens)))
             .collect()
     }
 
diff --git a/src/client/localai.rs b/src/client/localai.rs
index 291fc96..026ae27 100644
--- a/src/client/localai.rs
+++ b/src/client/localai.rs
@@ -1,4 +1,4 @@
-use super::openai::openai_build_body;
+use super::openai::{openai_build_body, openai_tokens_formula};
 use super::{ExtraConfig, LocalAIClient, ModelInfo, PromptKind, PromptType, SendData};
 
 use anyhow::Result;
@@ -45,7 +45,7 @@ impl LocalAIClient {
         local_config
             .models
             .iter()
-            .map(|v| ModelInfo::new(client, &v.name, v.max_tokens, index))
+            .map(|v| openai_tokens_formula(ModelInfo::new(index, client, &v.name).set_max_tokens(v.max_tokens)))
             .collect()
     }
 
diff --git a/src/client/openai.rs b/src/client/openai.rs
index b554f35..387baf0 100644
--- a/src/client/openai.rs
+++ b/src/client/openai.rs
@@ -38,7 +38,7 @@ impl OpenAIClient {
         let client = Self::name(local_config);
         MODELS
             .into_iter()
-            .map(|(name, max_tokens)| ModelInfo::new(client, name, Some(max_tokens), index))
+            .map(|(name, max_tokens)| openai_tokens_formula(ModelInfo::new(index, client, name).set_max_tokens(Some(max_tokens))))
             .collect()
     }
 
@@ -135,3 +135,7 @@ 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)
+}