diff options
| author | sigoden <sigoden@gmail.com> | 2023-11-01 22:15:55 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2023-11-01 22:15:55 +0800 |
| commit | f6da06dad9b2a76016209a7d58ad923c2c72f150 (patch) | |
| tree | 1c74f306c194d3289fe06ed31a39a00c27ad0ca9 /src/client | |
| parent | da3c541b681b140feb52b4297167ca233535fa90 (diff) | |
| download | aichat-f6da06dad9b2a76016209a7d58ad923c2c72f150.tar.gz | |
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
Diffstat (limited to 'src/client')
| -rw-r--r-- | src/client/azure_openai.rs | 4 | ||||
| -rw-r--r-- | src/client/localai.rs | 4 | ||||
| -rw-r--r-- | src/client/openai.rs | 6 |
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) +} |
