diff options
| -rw-r--r-- | README.md | 1 | ||||
| -rw-r--r-- | config.example.yaml | 6 | ||||
| -rw-r--r-- | src/client/mod.rs | 3 | ||||
| -rw-r--r-- | src/client/moonshot.rs | 67 |
4 files changed, 75 insertions, 2 deletions
@@ -53,6 +53,7 @@ Download it from [GitHub Releases](https://github.com/sigoden/aichat/releases), - VertexAI: Gemini-1/Gemini-1.5 (paid, vision) - Ernie (paid) - Qianwen (paid, vision) + - Moonshot (paid) - Support [Command Mode](#command) and [Chat-REPL Mode](#chat-repl) - Support [roles](#roles) - Support sessions (context-aware conversation) diff --git a/config.example.yaml b/config.example.yaml index 2d5cba0..8f8e22a 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -65,7 +65,7 @@ clients: api_key: Basic xxx # Set authorization header chat_endpoint: /chat # Optional field models: - - name: mistral + - name: llama2 max_input_tokens: 8192 # See https://learn.microsoft.com/en-us/azure/ai-services/openai/chatgpt-quickstart @@ -92,3 +92,7 @@ clients: # See https://help.aliyun.com/zh/dashscope/ - type: qianwen api_key: sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx + + # See https://platform.moonshot.cn/docs/intro + - type: moonshot + api_key: xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
\ No newline at end of file diff --git a/src/client/mod.rs b/src/client/mod.rs index b1d4979..7c7cfba 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -20,7 +20,8 @@ register_client!( AzureOpenAIConfig, AzureOpenAIClient ), + (vertexai, "vertexai", VertexAIConfig, VertexAIClient), (ernie, "ernie", ErnieConfig, ErnieClient), (qianwen, "qianwen", QianwenConfig, QianwenClient), - (vertexai, "vertexai", VertexAIConfig, VertexAIClient), + (moonshot, "moonshot", MoonshotConfig, MoonshotClient), ); diff --git a/src/client/moonshot.rs b/src/client/moonshot.rs new file mode 100644 index 0000000..a105e01 --- /dev/null +++ b/src/client/moonshot.rs @@ -0,0 +1,67 @@ +use super::openai::{openai_build_body, OPENAI_TOKENS_COUNT_FACTORS}; +use super::{ExtraConfig, MoonshotClient, Model, PromptType, SendData}; + +use crate::utils::PromptKind; + +use anyhow::Result; +use async_trait::async_trait; +use reqwest::{Client as ReqwestClient, RequestBuilder}; +use serde::Deserialize; + +const API_URL: &str = "https://api.moonshot.cn/v1/chat/completions"; + +const MODELS: [(&str, usize, &str); 3] = [ + // https://platform.moonshot.cn/docs/intro + ("moonshot-v1-8k", 8000, "text"), + ("moonshot-v1-32k", 32000, "text"), + ("moonshot-v1-128k", 128000, "text"), +]; + + +#[derive(Debug, Clone, Deserialize)] +pub struct MoonshotConfig { + pub name: Option<String>, + pub api_key: Option<String>, + pub extra: Option<ExtraConfig>, +} + +openai_compatible_client!(MoonshotClient); + +impl MoonshotClient { + config_get_fn!(api_key, get_api_key); + + pub const PROMPTS: [PromptType<'static>; 1] = [ + ("api_key", "API Key:", false, PromptKind::String), + ]; + + pub fn list_models(local_config: &MoonshotConfig) -> Vec<Model> { + 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)) + .set_tokens_count_factors(OPENAI_TOKENS_COUNT_FACTORS) + }) + .collect() + } + + 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()); + self.model.merge_extra_fields(&mut body); + + let url = API_URL; + + debug!("Moonshot Request: {url} {body}"); + + let mut builder = client.post(url).json(&body); + if let Some(api_key) = api_key { + builder = builder.bearer_auth(api_key); + } + + Ok(builder) + } +} |