From f6da06dad9b2a76016209a7d58ad923c2c72f150 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Wed, 1 Nov 2023 22:15:55 +0800
Subject: refactor: improve code quanity (#194)

- extends ModelInfo for tokens calculating
- refactor config/session.rs, improve export, render, getter/setter
- modify main.rs, allow --model override session.model
---
 src/client/azure_openai.rs | 4 ++--
 src/client/localai.rs      | 4 ++--
 src/client/openai.rs       | 6 +++++-
 3 files changed, 9 insertions(+), 5 deletions(-)

(limited to 'src/client')

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)
+}
-- 
cgit v1.2.3