summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
Diffstat
-rw-r--r--config.example.yaml9
-rw-r--r--src/client/azure_openai.rs15
-rw-r--r--src/client/claude.rs31
-rw-r--r--src/client/cohere.rs10
-rw-r--r--src/client/ernie.rs60
-rw-r--r--src/client/gemini.rs4
-rw-r--r--src/client/mistral.rs2
-rw-r--r--src/client/model.rs40
-rw-r--r--src/client/ollama.rs31
-rw-r--r--src/client/openai.rs13
-rw-r--r--src/client/openai_compatible.rs16
-rw-r--r--src/client/qianwen.rs10
-rw-r--r--src/client/vertexai.rs16
13 files changed, 160 insertions, 97 deletions
diff --git a/config.example.yaml b/config.example.yaml
index e3829a1..deed762 100644
--- a/config.example.yaml
+++ b/config.example.yaml
@@ -28,7 +28,14 @@ right_prompt: '{color.purple}{?session {?consume_tokens {consume_tokens}({consum
clients:
# All clients have the following configuration:
# - type: xxxx
- # name: nova # Only use it to distinguish clients with the same client type. Optional
+ # name: xxxx # Only use it to distinguish clients with the same client type. Optional
+ # models:
+ # - name: xxxx # The model name
+ # max_input_tokens: 100000 # Optional field
+ # max_output_tokens: 4096 # Optional field
+ # capabilities: text,vision # Optional field, supported capabilities: text, vision
+ # extra_fields: # Optional field, set custom parameters, will merge with the body json
+ # key: value
# extra:
# proxy: socks5://127.0.0.1:1080 # Specify https/socks5 proxy server. Note HTTPS_PROXY/ALL_PROXY also works.
# connect_timeout: 10 # Set a timeout in seconds for connect to server
diff --git a/src/client/azure_openai.rs b/src/client/azure_openai.rs
index 3d0cf3f..1726bbe 100644
--- a/src/client/azure_openai.rs
+++ b/src/client/azure_openai.rs
@@ -1,5 +1,5 @@
use super::openai::openai_build_body;
-use super::{AzureOpenAIClient, ExtraConfig, Model, ModelConfig, PromptType, SendData};
+use super::{convert_models, AzureOpenAIClient, ExtraConfig, Model, ModelConfig, PromptType, SendData};
use crate::utils::PromptKind;
@@ -37,23 +37,14 @@ impl AzureOpenAIClient {
pub fn list_models(local_config: &AzureOpenAIConfig) -> Vec<Model> {
let client_name = Self::name(local_config);
-
- local_config
- .models
- .iter()
- .map(|v| {
- Model::new(client_name, &v.name)
- .set_max_input_tokens(v.max_input_tokens)
- .set_capabilities(v.capabilities)
- })
- .collect()
+ convert_models(client_name, &local_config.models)
}
fn request_builder(&self, client: &ReqwestClient, data: SendData) -> Result<RequestBuilder> {
let api_base = self.get_api_base()?;
let api_key = self.get_api_key()?;
- let mut body = openai_build_body(data, self.model.name.clone());
+ let mut body = openai_build_body(data, &self.model);
self.model.merge_extra_fields(&mut body);
let url = format!(
diff --git a/src/client/claude.rs b/src/client/claude.rs
index 4da5128..d87e64a 100644
--- a/src/client/claude.rs
+++ b/src/client/claude.rs
@@ -15,11 +15,11 @@ use serde_json::{json, Value};
const API_BASE: &str = "https://api.anthropic.com/v1/messages";
-const MODELS: [(&str, usize, &str); 3] = [
+const MODELS: [(&str, usize, isize, &str); 3] = [
// https://docs.anthropic.com/claude/docs/models-overview
- ("claude-3-opus-20240229", 200000, "text,vision"),
- ("claude-3-sonnet-20240229", 200000, "text,vision"),
- ("claude-3-haiku-20240307", 200000, "text,vision"),
+ ("claude-3-opus-20240229", 200000, 4096, "text,vision"),
+ ("claude-3-sonnet-20240229", 200000, 4096, "text,vision"),
+ ("claude-3-haiku-20240307", 200000, 4096, "text,vision"),
];
#[derive(Debug, Clone, Deserialize)]
@@ -59,18 +59,21 @@ impl ClaudeClient {
let client_name = Self::name(local_config);
MODELS
.into_iter()
- .map(|(name, max_input_tokens, capabilities)| {
- Model::new(client_name, name)
- .set_capabilities(capabilities.into())
- .set_max_input_tokens(Some(max_input_tokens))
- })
+ .map(
+ |(name, max_input_tokens, max_output_tokens, capabilities)| {
+ Model::new(client_name, name)
+ .set_capabilities(capabilities.into())
+ .set_max_input_tokens(Some(max_input_tokens))
+ .set_max_output_tokens(Some(max_output_tokens))
+ },
+ )
.collect()
}
fn request_builder(&self, client: &ReqwestClient, data: SendData) -> Result<RequestBuilder> {
let api_key = self.get_api_key().ok();
- let body = build_body(data, self.model.name.clone())?;
+ let body = build_body(data, &self.model)?;
let url = API_BASE;
@@ -143,7 +146,7 @@ async fn send_message_streaming(builder: RequestBuilder, handler: &mut ReplyHand
Ok(())
}
-fn build_body(data: SendData, model: String) -> Result<Value> {
+fn build_body(data: SendData, model: &Model) -> Result<Value> {
let SendData {
mut messages,
temperature,
@@ -197,9 +200,11 @@ fn build_body(data: SendData, model: String) -> Result<Value> {
);
}
+ let max_tokens = model.max_output_tokens.unwrap_or(4096);
+
let mut body = json!({
- "model": model,
- "max_tokens": 4096,
+ "model": &model.name,
+ "max_tokens": max_tokens,
"messages": messages,
});
diff --git a/src/client/cohere.rs b/src/client/cohere.rs
index a92e238..bfea105 100644
--- a/src/client/cohere.rs
+++ b/src/client/cohere.rs
@@ -67,7 +67,7 @@ impl CohereClient {
fn request_builder(&self, client: &ReqwestClient, data: SendData) -> Result<RequestBuilder> {
let api_key = self.get_api_key().ok();
- let body = build_body(data, self.model.name.clone())?;
+ let body = build_body(data, &self.model)?;
let url = API_URL;
@@ -131,7 +131,7 @@ fn check_error(data: &Value) -> Result<()> {
}
}
-pub(crate) fn build_body(data: SendData, model: String) -> Result<Value> {
+pub(crate) fn build_body(data: SendData, model: &Model) -> Result<Value> {
let SendData {
mut messages,
temperature,
@@ -179,9 +179,13 @@ pub(crate) fn build_body(data: SendData, model: String) -> Result<Value> {
let message = message["message"].as_str().unwrap_or_default();
let mut body = json!({
- "model": model,
+ "model": &model.name,
"message": message,
});
+
+ if let Some(max_tokens) = model.max_output_tokens {
+ body["max_tokens"] = max_tokens.into();
+ }
if !messages.is_empty() {
body["chat_history"] = messages.into();
diff --git a/src/client/ernie.rs b/src/client/ernie.rs
index db6a969..0b10022 100644
--- a/src/client/ernie.rs
+++ b/src/client/ernie.rs
@@ -18,27 +18,50 @@ use std::{env, sync::Mutex};
const API_BASE: &str = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1";
const ACCESS_TOKEN_URL: &str = "https://aip.baidubce.com/oauth/2.0/token";
-const MODELS: [(&str, usize, &str); 7] = [
+const MODELS: [(&str, &str, usize, isize); 7] = [
// https://cloud.baidu.com/doc/WENXINWORKSHOP/s/clntwmv7t
- ("ernie-4.0-8k", 5120, "/wenxinworkshop/chat/completions_pro"),
(
- "ernie-3.5-8k",
+ "ernie-4.0-8k",
+ "/wenxinworkshop/chat/completions_pro",
5120,
+ 2048,
+ ),
+ (
+ "ernie-3.5-8k",
"/wenxinworkshop/chat/ernie-3.5-8k-0205",
+ 5120,
+ 2048,
),
(
"ernie-3.5-4k",
- 2048,
"/wenxinworkshop/chat/ernie-3.5-4k-0205",
+ 2048,
+ 2048,
+ ),
+ (
+ "ernie-speed-8k",
+ "/wenxinworkshop/chat/ernie_speed",
+ 7168,
+ 2048,
),
- ("ernie-speed-8k", 7168, "/wenxinworkshop/chat/ernie_speed"),
(
"ernie-speed-128k",
- 124000,
"/wenxinworkshop/chat/ernie-speed-128k",
+ 124000,
+ 4096,
+ ),
+ (
+ "ernie-lite-8k",
+ "/wenxinworkshop/chat/ernie-lite-8k",
+ 7168,
+ 2048,
+ ),
+ (
+ "ernie-tiny-8k",
+ "/wenxinworkshop/chat/ernie-tiny-8k",
+ 7168,
+ 2048,
),
- ("ernie-lite-8k", 7168, "/wenxinworkshop/chat/ernie-lite-8k"),
- ("ernie-tiny-8k", 7168, "/wenxinworkshop/chat/ernie-tiny-8k"),
];
lazy_static! {
@@ -85,17 +108,21 @@ impl ErnieClient {
let client_name = Self::name(local_config);
MODELS
.into_iter()
- .map(|(name, _, _)| Model::new(client_name, name)) // ERNIE tokenizer is different from cl100k_base
+ .map(|(name, _, max_input_tokens, max_output_tokens)| {
+ Model::new(client_name, name)
+ .set_max_input_tokens(Some(max_input_tokens))
+ .set_max_output_tokens(Some(max_output_tokens))
+ }) // ERNIE tokenizer is different from cl100k_base
.collect()
}
fn request_builder(&self, client: &ReqwestClient, data: SendData) -> Result<RequestBuilder> {
- let body = build_body(data, self.model.name.clone());
+ let body = build_body(data, &self.model);
- let model = self.model.name.clone();
- let (_, _, chat_endpoint) = MODELS
+ let model = &self.model.name;
+ let (_, chat_endpoint, _, _) = MODELS
.iter()
- .find(|(v, _, _)| v == &model)
+ .find(|(v, _, _, _)| v == model)
.ok_or_else(|| anyhow!("Miss Model '{}'", self.model.id()))?;
let access_token = ACCESS_TOKEN
@@ -207,7 +234,7 @@ fn check_error(data: &Value) -> Result<()> {
Ok(())
}
-fn build_body(data: SendData, _model: String) -> Value {
+fn build_body(data: SendData, model: &Model) -> Value {
let SendData {
mut messages,
temperature,
@@ -223,6 +250,11 @@ fn build_body(data: SendData, _model: String) -> Value {
if let Some(temperature) = temperature {
body["temperature"] = temperature.into();
}
+
+ if let Some(max_output_tokens) = model.max_output_tokens {
+ body["max_output_tokens"] = max_output_tokens.into();
+ }
+
if stream {
body["stream"] = true.into();
}
diff --git a/src/client/gemini.rs b/src/client/gemini.rs
index 0c60bee..5e42cd5 100644
--- a/src/client/gemini.rs
+++ b/src/client/gemini.rs
@@ -73,9 +73,9 @@ impl GeminiClient {
let block_threshold = self.config.block_threshold.clone();
- let body = build_body(data, self.model.name.clone(), block_threshold)?;
+ let body = build_body(data, &self.model, block_threshold)?;
- let model = self.model.name.clone();
+ let model = &self.model.name;
let url = format!("{API_BASE}{}:{}?key={}", model, func, api_key);
diff --git a/src/client/mistral.rs b/src/client/mistral.rs
index 611c25c..4ff4787 100644
--- a/src/client/mistral.rs
+++ b/src/client/mistral.rs
@@ -51,7 +51,7 @@ impl MistralClient {
fn request_builder(&self, client: &ReqwestClient, data: SendData) -> Result<RequestBuilder> {
let api_key = self.get_api_key().ok();
- let body = openai_build_body(data, self.model.name.clone());
+ let body = openai_build_body(data, &self.model);
let url = API_URL;
diff --git a/src/client/model.rs b/src/client/model.rs
index 03877b9..53d1834 100644
--- a/src/client/model.rs
+++ b/src/client/model.rs
@@ -13,6 +13,7 @@ pub struct Model {
pub client_name: String,
pub name: String,
pub max_input_tokens: Option<usize>,
+ pub max_output_tokens: Option<isize>,
pub extra_fields: Option<serde_json::Map<String, serde_json::Value>>,
pub capabilities: ModelCapabilities,
}
@@ -30,6 +31,7 @@ impl Model {
name: name.into(),
extra_fields: None,
max_input_tokens: None,
+ max_output_tokens: None,
capabilities: ModelCapabilities::Text,
}
}
@@ -90,6 +92,14 @@ impl Model {
self
}
+ pub fn set_max_output_tokens(mut self, max_output_tokens: Option<isize>) -> Self {
+ match max_output_tokens {
+ None | Some(0) => self.max_output_tokens = None,
+ _ => self.max_output_tokens = max_output_tokens,
+ }
+ self
+ }
+
pub fn messages_tokens(&self, messages: &[Message]) -> usize {
messages
.iter()
@@ -127,19 +137,43 @@ impl Model {
pub fn merge_extra_fields(&self, body: &mut serde_json::Value) {
if let (Some(body), Some(extra_fields)) = (body.as_object_mut(), &self.extra_fields) {
- for (k, v) in extra_fields {
- if !body.contains_key(k) {
- body.insert(k.clone(), v.clone());
+ for (key, extra_field) in extra_fields {
+ if body.contains_key(key) {
+ if let (Some(sub_body), Some(extra_field)) =
+ (body[key].as_object_mut(), extra_field.as_object())
+ {
+ for (subkey, sub_field) in extra_field {
+ if !sub_body.contains_key(subkey) {
+ sub_body.insert(subkey.clone(), sub_field.clone());
+ }
+ }
+ }
+ } else {
+ body.insert(key.clone(), extra_field.clone());
}
}
}
}
}
+pub fn convert_models(client_name: &str, models: &[ModelConfig]) -> Vec<Model> {
+ models
+ .iter()
+ .map(|v| {
+ Model::new(client_name, &v.name)
+ .set_capabilities(v.capabilities)
+ .set_max_input_tokens(v.max_input_tokens)
+ .set_max_output_tokens(v.max_output_tokens)
+ .set_extra_fields(v.extra_fields.clone())
+ })
+ .collect()
+}
+
#[derive(Debug, Clone, Deserialize)]
pub struct ModelConfig {
pub name: String,
pub max_input_tokens: Option<usize>,
+ pub max_output_tokens: Option<isize>,
pub extra_fields: Option<serde_json::Map<String, serde_json::Value>>,
#[serde(deserialize_with = "deserialize_capabilities")]
#[serde(default = "default_capabilities")]
diff --git a/src/client/ollama.rs b/src/client/ollama.rs
index e652634..1692942 100644
--- a/src/client/ollama.rs
+++ b/src/client/ollama.rs
@@ -1,6 +1,6 @@
use super::{
- message::*, patch_system_message, Client, ExtraConfig, Model, ModelConfig, OllamaClient,
- PromptType, ReplyHandler, SendData,
+ convert_models, message::*, patch_system_message, Client, ExtraConfig, Model, ModelConfig,
+ OllamaClient, PromptType, ReplyHandler, SendData,
};
use crate::utils::PromptKind;
@@ -59,23 +59,13 @@ impl OllamaClient {
pub fn list_models(local_config: &OllamaConfig) -> Vec<Model> {
let client_name = Self::name(local_config);
-
- local_config
- .models
- .iter()
- .map(|v| {
- Model::new(client_name, &v.name)
- .set_capabilities(v.capabilities)
- .set_max_input_tokens(v.max_input_tokens)
- .set_extra_fields(v.extra_fields.clone())
- })
- .collect()
+ convert_models(client_name, &local_config.models)
}
fn request_builder(&self, client: &ReqwestClient, data: SendData) -> Result<RequestBuilder> {
let api_key = self.get_api_key().ok();
- let mut body = build_body(data, self.model.name.clone())?;
+ let mut body = build_body(data, &self.model)?;
self.model.merge_extra_fields(&mut body);
let chat_endpoint = self.config.chat_endpoint.as_deref().unwrap_or("/api/chat");
@@ -133,7 +123,7 @@ async fn send_message_streaming(builder: RequestBuilder, handler: &mut ReplyHand
Ok(())
}
-fn build_body(data: SendData, model: String) -> Result<Value> {
+fn build_body(data: SendData, model: &Model) -> Result<Value> {
let SendData {
mut messages,
temperature,
@@ -189,15 +179,18 @@ fn build_body(data: SendData, model: String) -> Result<Value> {
}
let mut body = json!({
- "model": model,
+ "model": &model.name,
"messages": messages,
"stream": stream,
+ "options": {},
});
+ if let Some(num_predict) = model.max_output_tokens {
+ body["options"]["num_predict"] = num_predict.into();
+ }
+
if let Some(temperature) = temperature {
- body["options"] = json!({
- "temperature": temperature,
- });
+ body["options"]["temperature"] = temperature.into();
}
Ok(body)
diff --git a/src/client/openai.rs b/src/client/openai.rs
index 797ee98..535e739 100644
--- a/src/client/openai.rs
+++ b/src/client/openai.rs
@@ -58,7 +58,7 @@ impl OpenAIClient {
let api_key = self.get_api_key()?;
let api_base = self.get_api_base().unwrap_or_else(|_| API_BASE.to_string());
- let body = openai_build_body(data, self.model.name.clone());
+ let body = openai_build_body(data, &self.model);
let url = format!("{api_base}/chat/completions");
@@ -139,7 +139,7 @@ pub async fn openai_send_message_streaming(
Ok(())
}
-pub fn openai_build_body(data: SendData, model: String) -> Value {
+pub fn openai_build_body(data: SendData, model: &Model) -> Value {
let SendData {
messages,
temperature,
@@ -147,15 +147,16 @@ pub fn openai_build_body(data: SendData, model: String) -> Value {
} = data;
let mut body = json!({
- "model": model,
+ "model": &model.name,
"messages": messages,
});
- // The default max_tokens of gpt-4-vision-preview is only 16, we need to make it larger
- if model == "gpt-4-vision-preview" {
+ if let Some(max_tokens) = model.max_output_tokens {
+ body["max_tokens"] = json!(max_tokens);
+ } else if model.name == "gpt-4-vision-preview" {
+ // The default max_tokens of gpt-4-vision-preview is only 16, we need to make it larger
body["max_tokens"] = json!(4096);
}
-
if let Some(v) = temperature {
body["temperature"] = v.into();
}
diff --git a/src/client/openai_compatible.rs b/src/client/openai_compatible.rs
index 595b2cd..f25d931 100644
--- a/src/client/openai_compatible.rs
+++ b/src/client/openai_compatible.rs
@@ -1,5 +1,5 @@
use super::openai::openai_build_body;
-use super::{ExtraConfig, Model, ModelConfig, OpenAICompatibleClient, PromptType, SendData};
+use super::{convert_models, ExtraConfig, Model, ModelConfig, OpenAICompatibleClient, PromptType, SendData};
use crate::utils::PromptKind;
@@ -38,23 +38,13 @@ impl OpenAICompatibleClient {
pub fn list_models(local_config: &OpenAICompatibleConfig) -> Vec<Model> {
let client_name = Self::name(local_config);
-
- local_config
- .models
- .iter()
- .map(|v| {
- Model::new(client_name, &v.name)
- .set_capabilities(v.capabilities)
- .set_max_input_tokens(v.max_input_tokens)
- .set_extra_fields(v.extra_fields.clone())
- })
- .collect()
+ convert_models(client_name, &local_config.models)
}
fn request_builder(&self, client: &ReqwestClient, data: SendData) -> Result<RequestBuilder> {
let api_key = self.get_api_key().ok();
- let mut body = openai_build_body(data, self.model.name.clone());
+ let mut body = openai_build_body(data, &self.model);
self.model.merge_extra_fields(&mut body);
let chat_endpoint = self
diff --git a/src/client/qianwen.rs b/src/client/qianwen.rs
index 2034736..6225b96 100644
--- a/src/client/qianwen.rs
+++ b/src/client/qianwen.rs
@@ -97,7 +97,7 @@ impl QianwenClient {
true => API_URL_VL,
false => API_URL,
};
- let (body, has_upload) = build_body(data, self.model.name.clone(), is_vl)?;
+ let (body, has_upload) = build_body(data, &self.model, is_vl)?;
debug!("Qianwen Request: {url} {body}");
@@ -180,7 +180,7 @@ fn check_error(data: &Value) -> Result<()> {
Ok(())
}
-fn build_body(data: SendData, model: String, is_vl: bool) -> Result<(Value, bool)> {
+fn build_body(data: SendData, model: &Model, is_vl: bool) -> Result<(Value, bool)> {
let SendData {
messages,
temperature,
@@ -233,6 +233,10 @@ fn build_body(data: SendData, model: String, is_vl: bool) -> Result<(Value, bool
parameters["incremental_output"] = true.into();
}
+ if let Some(max_tokens) = model.max_output_tokens {
+ parameters["max_tokens"] = max_tokens.into();
+ }
+
if let Some(v) = temperature {
parameters["temperature"] = v.into();
}
@@ -240,7 +244,7 @@ fn build_body(data: SendData, model: String, is_vl: bool) -> Result<(Value, bool
};
let body = json!({
- "model": model,
+ "model": &model.name,
"input": input,
"parameters": parameters
});
diff --git a/src/client/vertexai.rs b/src/client/vertexai.rs
index 88035ec..eceeb2c 100644
--- a/src/client/vertexai.rs
+++ b/src/client/vertexai.rs
@@ -81,9 +81,9 @@ impl VertexAIClient {
let block_threshold = self.config.block_threshold.clone();
- let body = build_body(data, self.model.name.clone(), block_threshold)?;
+ let body = build_body(data, &self.model, block_threshold)?;
- let model = self.model.name.clone();
+ let model = &self.model.name;
let url = format!("{api_base}/{}:{}", model, func);
@@ -176,7 +176,7 @@ fn check_error(data: &Value) -> Result<()> {
pub(crate) fn build_body(
data: SendData,
- _model: String,
+ model: &Model,
block_threshold: Option<String>,
) -> Result<Value> {
let SendData {
@@ -228,7 +228,7 @@ pub(crate) fn build_body(
);
}
- let mut body = json!({ "contents": contents });
+ let mut body = json!({ "contents": contents, "generationConfig": {} });
if let Some(block_threshold) = block_threshold {
body["safetySettings"] = json!([
@@ -239,10 +239,12 @@ pub(crate) fn build_body(
]);
}
+ if let Some(max_output_tokens) = model.max_output_tokens {
+ body["generationConfig"]["maxOutputTokens"] = max_output_tokens.into();
+ }
+
if let Some(temperature) = temperature {
- body["generationConfig"] = json!({
- "temperature": temperature,
- });
+ body["generationConfig"]["temperature"] = temperature.into();
}
Ok(body)
generated by cgit v1.2.3 (git 2.39.1) at 2026-10-03 18:10:55 +0000