mirror of https://github.com/sigoden/aichat
feat: support moonshot (#369)
parent
527da63d18
commit
774d991144
@ -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)
|
||||||
|
}
|
||||||
|
}
|
Loading…
Reference in New Issue