refactor: set tokens count factors (#195)

1 files changed, 6 insertions, 2 deletions

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()
     }