feat(google_drive): add image resizing logic from developer, and use Content::Image (#1735)

This commit is contained in:
Kalvin C
2025-03-18 08:32:27 -07:00
committed by GitHub
parent d6c0cc9dc9
commit ca8b38470b
+91 -10
View File
@@ -1,6 +1,8 @@
mod oauth_pkce;
pub mod storage;
use anyhow::{Context, Error};
use base64::Engine;
use indoc::indoc;
use oauth_pkce::PkceOAuth2Client;
use regex::Regex;
@@ -632,6 +634,67 @@ impl GoogleDriveRouter {
image_regex.replace_all(input, "").to_string()
}
// Helper function that processes one captured image.
// It decodes the base64 data, resizes the image if its width exceeds `max_width`,
// and then returns a new image tag (always output as PNG).
// logic copied from developer/mod.rs
fn process_image(&self, caps: &regex::Captures, max_width: u32) -> Result<Content, Error> {
let base64_data = &caps["data"];
// Decode the Base64 data.
let image_bytes = base64::prelude::BASE64_STANDARD
.decode(base64_data)
.context("Failed to decode base64 image data")?;
// Load the image from the decoded bytes.
let img = xcap::image::load_from_memory(&image_bytes)
.context("Failed to load image from memory")?;
// Resize the image if necessary.
let mut processed_image = img;
if processed_image.width() > max_width {
let scale = max_width as f32 / processed_image.width() as f32;
let new_height = (processed_image.height() as f32 * scale) as u32;
processed_image = xcap::image::DynamicImage::ImageRgba8(xcap::image::imageops::resize(
&processed_image,
max_width,
new_height,
xcap::image::imageops::FilterType::Lanczos3,
));
}
// Write the processed image to an in-memory buffer in PNG format.
let mut buffer: Vec<u8> = Vec::new();
processed_image
.write_to(&mut Cursor::new(&mut buffer), xcap::image::ImageFormat::Png)
.context("Failed to write processed image to buffer")?;
// Re-encode the buffer back into a Base64 string.
let data = base64::prelude::BASE64_STANDARD.encode(&buffer);
Ok(Content::image(data, "image/png"))
}
/// Resizes all base64-encoded images found in the input string.
/// If any image fails to process, an error is returned.
fn resize_images(&self, input: &str) -> Result<Vec<Content>, Error> {
// Regex to match and capture the MIME type and Base64 data.
let image_regex =
Regex::new(r"<data:image/(?P<mime>[a-zA-Z0-9.+-]+);base64,(?P<data>[^>]+)>")
.context("Failed to compile regex")?;
let mut result: Vec<Content> = Vec::new();
// Iterate over all matches, process them, and rebuild the output string.
for caps in image_regex.captures_iter(input) {
let processed_tag = self
.process_image(&caps, 768)
.context("Failed to process one of the images")?;
result.push(processed_tag);
}
Ok(result)
}
// Downloading content with alt=media only works if the file is stored in Drive.
// To download Google Docs, Sheets, and Slides use files.export instead.
async fn export_google_file(
@@ -665,13 +728,22 @@ impl GoogleDriveRouter {
Ok(r) => {
if let Ok(body) = r.into_body().collect().await {
if let Ok(response) = String::from_utf8(body.to_bytes().to_vec()) {
let content = if !include_images {
self.strip_image_body(&response)
if !include_images {
let content = self.strip_image_body(&response);
Ok(vec![Content::text(content).with_priority(0.1)])
} else {
response
};
let images = self.resize_images(&response).map_err(|e| {
ToolError::ExecutionError(format!(
"Failed to resize image(s): {}",
e
))
})?;
Ok(vec![Content::text(content).with_priority(0.1)])
let content = self.strip_image_body(&response);
Ok(std::iter::once(Content::text(content).with_priority(0.1))
.chain(images.iter().cloned())
.collect::<Vec<Content>>())
}
} else {
Err(ToolError::ExecutionError(format!(
"Failed to export google drive to string, {}.",
@@ -717,13 +789,22 @@ impl GoogleDriveRouter {
if mime_type.starts_with("text/") || mime_type == "application/json" {
if let Ok(body) = r.0.into_body().collect().await {
if let Ok(response) = String::from_utf8(body.to_bytes().to_vec()) {
let content = if !include_images {
self.strip_image_body(&response)
if !include_images {
let content = self.strip_image_body(&response);
Ok(vec![Content::text(content).with_priority(0.1)])
} else {
response
};
let images = self.resize_images(&response).map_err(|e| {
ToolError::ExecutionError(format!(
"Failed to resize image(s): {}",
e
))
})?;
Ok(vec![Content::text(content).with_priority(0.1)])
let content = self.strip_image_body(&response);
Ok(std::iter::once(Content::text(content).with_priority(0.1))
.chain(images.iter().cloned())
.collect::<Vec<Content>>())
}
} else {
Err(ToolError::ExecutionError(format!(
"Failed to convert google drive to string, {}.",