mirror of https://github.com/sigoden/aichat
feat: support gemini (#273)
sigoden
6 months ago
committed by
GitHub
parent
3adfeb1ae7
commit
6286251d32
@ -0,0 +1,241 @@
|
||||
use super::{
|
||||
patch_system_message, Client, ExtraConfig, GeminiClient, Model, PromptType, SendData,
|
||||
TokensCountFactors,
|
||||
};
|
||||
|
||||
use crate::{client::*, config::GlobalConfig, render::ReplyHandler, utils::PromptKind};
|
||||
|
||||
use anyhow::{anyhow, bail, Result};
|
||||
use async_trait::async_trait;
|
||||
use futures_util::StreamExt;
|
||||
use reqwest::{Client as ReqwestClient, RequestBuilder};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
const API_BASE: &str = "https://generativelanguage.googleapis.com/v1beta/models/";
|
||||
|
||||
const MODELS: [(&str, usize); 3] = [
|
||||
("gemini-pro", 32768),
|
||||
("gemini-pro-vision", 16384),
|
||||
("gemini-ultra", 32768),
|
||||
];
|
||||
|
||||
const TOKENS_COUNT_FACTORS: TokensCountFactors = (5, 2);
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Default)]
|
||||
pub struct GeminiConfig {
|
||||
pub name: Option<String>,
|
||||
pub api_key: Option<String>,
|
||||
pub extra: Option<ExtraConfig>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Client for GeminiClient {
|
||||
fn config(&self) -> (&GlobalConfig, &Option<ExtraConfig>) {
|
||||
(&self.global_config, &self.config.extra)
|
||||
}
|
||||
|
||||
async fn send_message_inner(&self, client: &ReqwestClient, data: SendData) -> Result<String> {
|
||||
let builder = self.request_builder(client, data)?;
|
||||
send_message(builder).await
|
||||
}
|
||||
|
||||
async fn send_message_streaming_inner(
|
||||
&self,
|
||||
client: &ReqwestClient,
|
||||
handler: &mut ReplyHandler,
|
||||
data: SendData,
|
||||
) -> Result<()> {
|
||||
let builder = self.request_builder(client, data)?;
|
||||
send_message_streaming(builder, handler).await
|
||||
}
|
||||
}
|
||||
|
||||
impl GeminiClient {
|
||||
config_get_fn!(api_key, get_api_key);
|
||||
|
||||
pub const PROMPTS: [PromptType<'static>; 1] =
|
||||
[("api_key", "API Key:", true, PromptKind::String)];
|
||||
|
||||
pub fn list_models(local_config: &GeminiConfig) -> Vec<Model> {
|
||||
let client_name = Self::name(local_config);
|
||||
MODELS
|
||||
.into_iter()
|
||||
.map(|(name, max_tokens)| {
|
||||
Model::new(client_name, name)
|
||||
.set_max_tokens(Some(max_tokens))
|
||||
.set_tokens_count_factors(TOKENS_COUNT_FACTORS)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn request_builder(&self, client: &ReqwestClient, data: SendData) -> Result<RequestBuilder> {
|
||||
let api_key = self.get_api_key()?;
|
||||
|
||||
let func = match data.stream {
|
||||
true => "streamGenerateContent",
|
||||
false => "generateContent",
|
||||
};
|
||||
|
||||
let body = build_body(data, self.model.name.clone())?;
|
||||
|
||||
let model = self.model.name.clone();
|
||||
|
||||
let url = format!("{API_BASE}{}:{}?key={}", model, func, api_key);
|
||||
|
||||
debug!("Gemini Request: {url} {body}");
|
||||
|
||||
let builder = client.post(url).json(&body);
|
||||
|
||||
Ok(builder)
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_message(builder: RequestBuilder) -> Result<String> {
|
||||
let res = builder.send().await?;
|
||||
let status = res.status();
|
||||
let data: Value = res.json().await?;
|
||||
if status != 200 {
|
||||
check_error(&data)?;
|
||||
}
|
||||
let output = data["candidates"][0]["content"]["parts"][0]["text"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow!("Invalid response data: {data}"))?;
|
||||
Ok(output.to_string())
|
||||
}
|
||||
|
||||
async fn send_message_streaming(builder: RequestBuilder, handler: &mut ReplyHandler) -> Result<()> {
|
||||
let res = builder.send().await?;
|
||||
if res.status() != 200 {
|
||||
let data: Value = res.json().await?;
|
||||
check_error(&data)?;
|
||||
} else {
|
||||
let mut buffer = vec![];
|
||||
let mut cursor = 0;
|
||||
let mut start = 0;
|
||||
let mut balances = vec![];
|
||||
let mut quoting = false;
|
||||
let mut stream = res.bytes_stream();
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk?;
|
||||
let chunk = std::str::from_utf8(&chunk)?;
|
||||
buffer.extend(chunk.chars());
|
||||
for i in cursor..buffer.len() {
|
||||
let ch = buffer[i];
|
||||
if quoting {
|
||||
if ch == '"' && buffer[i-1] != '\\' {
|
||||
quoting = false;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
match ch {
|
||||
'"' => quoting = true,
|
||||
'{' => {
|
||||
if balances.is_empty() {
|
||||
start = i;
|
||||
}
|
||||
balances.push(ch);
|
||||
}
|
||||
'[' => {
|
||||
if start != 0 {
|
||||
balances.push(ch);
|
||||
}
|
||||
}
|
||||
'}' => {
|
||||
balances.pop();
|
||||
if balances.is_empty() {
|
||||
let value: String = buffer[start..=i].iter().collect();
|
||||
let value: Value = serde_json::from_str(&value)?;
|
||||
if let Some(text) =
|
||||
value["candidates"][0]["content"]["parts"][0]["text"].as_str()
|
||||
{
|
||||
handler.text(text)?;
|
||||
} else {
|
||||
bail!("Invalid response data: {value}")
|
||||
}
|
||||
}
|
||||
}
|
||||
']' => {
|
||||
balances.pop();
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
cursor = buffer.len();
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn check_error(data: &Value) -> Result<()> {
|
||||
if let Some((Some(status), Some(message))) = data[0]["error"].as_object().map(|v| {
|
||||
(
|
||||
v.get("status").and_then(|v| v.as_str()),
|
||||
v.get("message").and_then(|v| v.as_str()),
|
||||
)
|
||||
}) {
|
||||
bail!("{status}: {message}")
|
||||
} else {
|
||||
bail!("Error {}", data);
|
||||
}
|
||||
}
|
||||
|
||||
fn build_body(data: SendData, _model: String) -> Result<Value> {
|
||||
let SendData {
|
||||
mut messages,
|
||||
temperature,
|
||||
..
|
||||
} = data;
|
||||
|
||||
patch_system_message(&mut messages);
|
||||
|
||||
let mut invalid_urls = vec![];
|
||||
let contents: Vec<Value> = messages
|
||||
.into_iter()
|
||||
.map(|message| {
|
||||
let role = match message.role {
|
||||
MessageRole::User => "user",
|
||||
_ => "model",
|
||||
};
|
||||
match message.content {
|
||||
MessageContent::Text(text) => json!({
|
||||
"role": role,
|
||||
"parts": [{ "text": text }]
|
||||
}),
|
||||
MessageContent::Array(list) => {
|
||||
let list: Vec<Value> = list
|
||||
.into_iter()
|
||||
.map(|item| match item {
|
||||
MessageContentPart::Text { text } => json!({"text": text}),
|
||||
MessageContentPart::ImageUrl { image_url: ImageUrl { url } } => {
|
||||
if let Some((mime_type, data)) = url.strip_prefix("data:").and_then(|v| v.split_once(";base64,")) {
|
||||
json!({ "inline_data": { "mime_type": mime_type, "data": data } })
|
||||
} else {
|
||||
invalid_urls.push(url.clone());
|
||||
json!({ "url": url })
|
||||
}
|
||||
},
|
||||
})
|
||||
.collect();
|
||||
json!({ "role": role, "parts": list })
|
||||
}
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
if !invalid_urls.is_empty() {
|
||||
bail!("The model does not support non-data URLs: {:?}", invalid_urls);
|
||||
}
|
||||
|
||||
let mut body = json!({
|
||||
"contents": contents,
|
||||
});
|
||||
|
||||
if let Some(temperature) = temperature {
|
||||
body["generationConfig"] = json!({
|
||||
"temperature": temperature,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(body)
|
||||
}
|
||||
Loading…
Reference in New Issue