diff options
Diffstat (limited to 'src/client/qianwen.rs')
| -rw-r--r-- | src/client/qianwen.rs | 147 |
1 files changed, 128 insertions, 19 deletions
diff --git a/src/client/qianwen.rs b/src/client/qianwen.rs index ed4b6e6..022e286 100644 --- a/src/client/qianwen.rs +++ b/src/client/qianwen.rs @@ -1,14 +1,23 @@ use super::{message::*, Client, ExtraConfig, Model, PromptType, QianwenClient, SendData}; -use crate::{config::GlobalConfig, render::ReplyHandler, utils::PromptKind}; +use crate::{ + config::GlobalConfig, + render::ReplyHandler, + utils::{sha256sum, PromptKind}, +}; -use anyhow::{anyhow, bail, Result}; +use anyhow::{anyhow, bail, Context, Result}; use async_trait::async_trait; +use base64::{engine::general_purpose::STANDARD, Engine}; use futures_util::StreamExt; -use reqwest::{Client as ReqwestClient, RequestBuilder}; +use reqwest::{ + multipart::{Form, Part}, + Client as ReqwestClient, RequestBuilder, +}; use reqwest_eventsource::{Error as EventSourceError, Event, RequestBuilderExt}; use serde::Deserialize; use serde_json::{json, Value}; +use std::borrow::BorrowMut; const API_URL: &str = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"; @@ -37,7 +46,13 @@ impl Client for QianwenClient { (&self.global_config, &self.config.extra) } - async fn send_message_inner(&self, client: &ReqwestClient, data: SendData) -> Result<String> { + async fn send_message_inner( + &self, + client: &ReqwestClient, + mut data: SendData, + ) -> Result<String> { + let api_key = self.get_api_key()?; + patch_messages(&self.model.name, &api_key, &mut data.messages).await?; let builder = self.request_builder(client, data)?; send_message(builder, self.is_vl()).await } @@ -46,8 +61,10 @@ impl Client for QianwenClient { &self, client: &ReqwestClient, handler: &mut ReplyHandler, - data: SendData, + mut data: SendData, ) -> Result<()> { + let api_key = self.get_api_key()?; + patch_messages(&self.model.name, &api_key, &mut data.messages).await?; let builder = self.request_builder(client, data)?; send_message_streaming(builder, handler, self.is_vl()).await } @@ -79,7 +96,7 @@ impl QianwenClient { true => API_URL_VL, false => API_URL, }; - let body = build_body(data, self.model.name.clone(), is_vl)?; + let (body, has_upload) = build_body(data, self.model.name.clone(), is_vl)?; debug!("Qianwen Request: {url} {body}"); @@ -87,6 +104,9 @@ impl QianwenClient { if stream { builder = builder.header("X-DashScope-SSE", "enable"); } + if has_upload { + builder = builder.header("X-DashScope-OssResourceResolve", "enable"); + } Ok(builder) } @@ -126,7 +146,8 @@ async fn send_message_streaming( let data: Value = serde_json::from_str(&message.data)?; check_error(&data)?; if is_vl { - let text = data["output"]["choices"][0]["message"]["content"][0]["text"].as_str(); + let text = + data["output"]["choices"][0]["message"]["content"][0]["text"].as_str(); if let Some(text) = text { let text = &text[offset..]; handler.text(text)?; @@ -158,17 +179,15 @@ fn check_error(data: &Value) -> Result<()> { Ok(()) } -fn build_body(data: SendData, model: String, is_vl: bool) -> Result<Value> { +fn build_body(data: SendData, model: String, is_vl: bool) -> Result<(Value, bool)> { let SendData { messages, temperature, stream, } = data; - + let mut has_upload = false; let (input, parameters) = if is_vl { - let mut exist_embeded_image = false; - let messages: Vec<Value> = messages .into_iter() .map(|message| { @@ -182,11 +201,11 @@ fn build_body(data: SendData, model: String, is_vl: bool) -> Result<Value> { MessageContentPart::ImageUrl { image_url: ImageUrl { url }, } => { - if url.starts_with("data:") { - exist_embeded_image = true; + if url.starts_with("oss:") { + has_upload = true; } json!({"image": url}) - }, + } }) .collect(), }; @@ -194,10 +213,6 @@ fn build_body(data: SendData, model: String, is_vl: bool) -> Result<Value> { }) .collect(); - if exist_embeded_image { - bail!("The model does not support embeded images"); - } - let input = json!({ "messages": messages, }); @@ -228,5 +243,99 @@ fn build_body(data: SendData, model: String, is_vl: bool) -> Result<Value> { "input": input, "parameters": parameters }); - Ok(body) + Ok((body, has_upload)) +} + +/// Patch messsages, upload emebeded images to oss +async fn patch_messages(model: &str, api_key: &str, messages: &mut Vec<Message>) -> Result<()> { + for message in messages { + if let MessageContent::Array(list) = message.content.borrow_mut() { + for item in list { + if let MessageContentPart::ImageUrl { + image_url: ImageUrl { url }, + } = item + { + if url.starts_with("data:") { + *url = upload(model, api_key, url) + .await + .with_context(|| "Failed to upload embeded image to oss")?; + } + } + } + } + } + Ok(()) +} + +#[derive(Debug, Deserialize)] +struct Policy { + data: PolicyData, +} + +#[derive(Debug, Deserialize)] +struct PolicyData { + policy: String, + signature: String, + upload_dir: String, + upload_host: String, + oss_access_key_id: String, + x_oss_object_acl: String, + x_oss_forbid_overwrite: String, +} + +/// Upload image to dashscope +async fn upload(model: &str, api_key: &str, url: &str) -> Result<String> { + let (mime_type, data) = url + .strip_prefix("data:") + .and_then(|v| v.split_once(";base64,")) + .ok_or_else(|| anyhow!("Invalid image url"))?; + let mut name = sha256sum(data); + if let Some(ext) = mime_type.strip_prefix("image/") { + name.push('.'); + name.push_str(ext); + } + let data = STANDARD.decode(data)?; + + let client = reqwest::Client::new(); + let policy: Policy = client + .get(format!( + "https://dashscope.aliyuncs.com/api/v1/uploads?action=getPolicy&model={model}" + )) + .header("Authorization", format!("Bearer {api_key}")) + .send() + .await? + .json() + .await?; + let PolicyData { + policy, + signature, + upload_dir, + upload_host, + oss_access_key_id, + x_oss_object_acl, + x_oss_forbid_overwrite, + .. + } = policy.data; + + let key = format!("{upload_dir}/{name}"); + let file = Part::bytes(data).file_name(name).mime_str(mime_type)?; + let form = Form::new() + .text("OSSAccessKeyId", oss_access_key_id) + .text("Signature", signature) + .text("policy", policy) + .text("key", key.clone()) + .text("x-oss-object-acl", x_oss_object_acl) + .text("x-oss-forbid-overwrite", x_oss_forbid_overwrite) + .text("success_action_status", "200") + .text("x-oss-content-type", mime_type.to_string()) + .part("file", file); + + let res = client.post(upload_host).multipart(form).send().await?; + + let status = res.status(); + if res.status() != 200 { + let text = res.text().await?; + bail!("{status}, {text}") + } + Ok(format!("oss://{key}")) } |