218 lines
10 KiB
Rust
218 lines
10 KiB
Rust
// Files and pictures the person gives a bot with a message (the original: 附件 — drag files in,
|
||
// paste a picture, or pick files; pictures go to the model, documents are read for it).
|
||
// Each one is copied into AnyBot's data folder, so the chat keeps working after the original moves.
|
||
|
||
use base64::Engine;
|
||
use serde::{Deserialize, Serialize};
|
||
use std::path::{Path, PathBuf};
|
||
|
||
/// Largest file accepted (estimated).
|
||
pub const MAX_FILE_BYTES: u64 = 52_428_800;
|
||
/// Text read from one file for the model (estimated; the original reads 20,000 characters of mail
|
||
/// and notes).
|
||
pub const MAX_TEXT_CHARS: usize = 20_000;
|
||
|
||
#[derive(Clone, Serialize, Deserialize, Debug, PartialEq)]
|
||
pub struct Attachment {
|
||
pub name: String,
|
||
/// The copy in AnyBot's data folder.
|
||
pub path: String,
|
||
#[serde(default)]
|
||
pub image: bool,
|
||
#[serde(default)]
|
||
pub bytes: u64,
|
||
/// Where it was also put on the bot's own computer (empty when not).
|
||
#[serde(default)]
|
||
pub on_box: String,
|
||
}
|
||
|
||
fn dir_for(bot_id: u64) -> PathBuf {
|
||
crate::model::State::path().parent().map(|p| p.to_path_buf()).unwrap_or_else(std::env::temp_dir).join("attachments").join(format!("bot-{bot_id}"))
|
||
}
|
||
|
||
pub fn is_image_name(name: &str) -> bool {
|
||
let n = name.to_lowercase();
|
||
[".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".heic", ".tif", ".tiff"].iter().any(|e| n.ends_with(e))
|
||
}
|
||
|
||
/// Copy a file in as an attachment.
|
||
pub fn store(bot_id: u64, src: &Path) -> Result<Attachment, String> {
|
||
let meta = std::fs::metadata(src).map_err(|e| format!("讀不到 {}：{e}", src.display()))?;
|
||
if meta.is_dir() {
|
||
return Err(format!("「{}」是資料夾，請壓縮後再附加，或請 Bot 直接處理那個資料夾。", src.display()));
|
||
}
|
||
if meta.len() > MAX_FILE_BYTES {
|
||
return Err(format!("「{}」太大了（最多 50 MB）", src.display()));
|
||
}
|
||
let name = src.file_name().map(|n| n.to_string_lossy().to_string()).unwrap_or_else(|| "附件".into());
|
||
let dir = dir_for(bot_id);
|
||
std::fs::create_dir_all(&dir).map_err(|e| e.to_string())?;
|
||
let dst = dir.join(format!("{}-{name}", chrono::Local::now().format("%Y%m%d%H%M%S%3f")));
|
||
std::fs::copy(src, &dst).map_err(|e| format!("複製失敗：{e}"))?;
|
||
Ok(Attachment { image: is_image_name(&name), name, path: dst.display().to_string(), bytes: meta.len(), on_box: String::new() })
|
||
}
|
||
|
||
/// A picture or file the bot hands over (a screen it shows, a file fetched from another computer).
|
||
pub fn store_bytes(bot_id: u64, name: &str, bytes: &[u8]) -> Result<Attachment, String> {
|
||
if bytes.len() as u64 > MAX_FILE_BYTES {
|
||
return Err(format!("「{name}」太大了（最多 50 MB）"));
|
||
}
|
||
let dir = dir_for(bot_id);
|
||
std::fs::create_dir_all(&dir).map_err(|e| e.to_string())?;
|
||
let dst = dir.join(format!("{}-{name}", chrono::Local::now().format("%Y%m%d%H%M%S%3f")));
|
||
std::fs::write(&dst, bytes).map_err(|e| format!("存不下來：{e}"))?;
|
||
Ok(Attachment { image: is_image_name(name), name: name.to_string(), path: dst.display().to_string(), bytes: bytes.len() as u64, on_box: String::new() })
|
||
}
|
||
|
||
/// A picture on the clipboard (e.g. a screenshot copied with ⌃⇧⌘4) becomes an attachment.
|
||
pub fn paste_image(bot_id: u64) -> Option<Attachment> {
|
||
let dir = dir_for(bot_id);
|
||
std::fs::create_dir_all(&dir).ok()?;
|
||
let name = format!("貼上的圖片 {}.png", chrono::Local::now().format("%H%M%S"));
|
||
let dst = dir.join(format!("{}-{name}", chrono::Local::now().format("%Y%m%d%H%M%S%3f")));
|
||
if !cfg!(target_os = "macos") {
|
||
let png = crate::platform::clipboard_png_linux()?;
|
||
std::fs::write(&dst, &png).ok()?;
|
||
return Some(Attachment { name, path: dst.display().to_string(), image: true, bytes: png.len() as u64, on_box: String::new() });
|
||
}
|
||
let script = r#"ObjC.import('AppKit');
|
||
function run(argv) {
|
||
const img = $.NSImage.alloc.initWithPasteboard($.NSPasteboard.generalPasteboard);
|
||
if (img.isNil()) return 'none';
|
||
const rep = $.NSBitmapImageRep.imageRepWithData(img.TIFFRepresentation);
|
||
const png = rep.representationUsingTypeProperties($.NSBitmapImageFileTypePNG, $({}));
|
||
return png.writeToFileAtomically(argv[0], true) ? 'ok' : 'fail';
|
||
}"#;
|
||
let out = std::process::Command::new("osascript").args(["-l", "JavaScript", "-e", script, "--"]).arg(&dst).output().ok()?;
|
||
if String::from_utf8_lossy(&out.stdout).trim() != "ok" {
|
||
return None;
|
||
}
|
||
let bytes = std::fs::metadata(&dst).map(|m| m.len()).unwrap_or(0);
|
||
Some(Attachment { name, path: dst.display().to_string(), image: true, bytes, on_box: String::new() })
|
||
}
|
||
|
||
/// The Finder file picker (several files at once).
|
||
pub fn pick_files() -> Vec<PathBuf> {
|
||
if !cfg!(target_os = "macos") {
|
||
return crate::platform::pick_files_linux(&crate::i18n::t("選擇要附加的檔案"));
|
||
}
|
||
let script = r#"function run() {
|
||
const app = Application.currentApplication(); app.includeStandardAdditions = true;
|
||
try { return app.chooseFile({withPrompt: '選擇要附加的檔案', multipleSelectionsAllowed: true}).map(f => f.toString()).join('\n'); } catch (e) { return ''; }
|
||
}"#;
|
||
std::process::Command::new("osascript")
|
||
.args(["-l", "JavaScript", "-e", script])
|
||
.output()
|
||
.map(|o| String::from_utf8_lossy(&o.stdout).lines().map(str::trim).filter(|l| !l.is_empty()).map(PathBuf::from).collect())
|
||
.unwrap_or_default()
|
||
}
|
||
|
||
pub fn human(b: u64) -> String {
|
||
crate::cloudbox::human(b)
|
||
}
|
||
|
||
/// A picture, sized like a screenshot, as base64 JPEG.
|
||
pub fn image_b64(path: &str) -> Option<String> {
|
||
// The viewer's decoder: every common format, SVG, and HEIC/AVIF through sips or ffmpeg.
|
||
let img = image::DynamicImage::ImageRgba8(crate::preview::decode_image(std::path::Path::new(path))?);
|
||
let img = img.to_rgb8();
|
||
let (w, h) = img.dimensions();
|
||
let (fw, fh, resize) = crate::computer::fit(w, h);
|
||
let img = if resize { image::imageops::resize(&img, fw, fh, image::imageops::FilterType::Triangle) } else { img };
|
||
crate::computer::encode_jpeg(&img).map(|j| base64::engine::general_purpose::STANDARD.encode(j))
|
||
}
|
||
|
||
fn clip(s: &str) -> String {
|
||
let n = s.chars().count();
|
||
if n <= MAX_TEXT_CHARS {
|
||
s.to_string()
|
||
} else {
|
||
format!("{}\n…（後面省略 {} 個字）", s.chars().take(MAX_TEXT_CHARS).collect::<String>(), n - MAX_TEXT_CHARS)
|
||
}
|
||
}
|
||
|
||
/// The text inside a document, when it can be read here.
|
||
pub fn read_text(path: &str, name: &str) -> Option<String> {
|
||
let n = name.to_lowercase();
|
||
let ext = n.rsplit('.').next().unwrap_or("");
|
||
// Word and web pages: read in AnyBot itself (same on every system).
|
||
if ext == "docx" {
|
||
if let Ok(md) = crate::preview::docx_markdown(std::path::Path::new(path)) {
|
||
return Some(md).filter(|t| !t.trim().is_empty());
|
||
}
|
||
}
|
||
if ext == "html" || ext == "htm" {
|
||
let bytes = std::fs::read(path).ok()?;
|
||
return Some(crate::preview::html_markdown(&String::from_utf8_lossy(&bytes))).filter(|t| !t.trim().is_empty());
|
||
}
|
||
if ["docx", "doc", "rtf", "rtfd", "odt", "webarchive", "pages"].contains(&ext) && cfg!(target_os = "macos") {
|
||
let out = std::process::Command::new("textutil").args(["-convert", "txt", "-stdout", path]).output().ok()?;
|
||
return out.status.success().then(|| String::from_utf8_lossy(&out.stdout).to_string()).filter(|t| !t.trim().is_empty());
|
||
}
|
||
if ext == "pdf" {
|
||
// pdftotext keeps the layout best; without it, read the text in AnyBot (pure Rust).
|
||
if let Some(t) = std::process::Command::new("pdftotext").env("PATH", crate::cloudbox::tool_path()).args(["-layout", path, "-"]).output().ok().filter(|o| o.status.success()).map(|o| String::from_utf8_lossy(&o.stdout).to_string()).filter(|t| !t.trim().is_empty()) {
|
||
return Some(t);
|
||
}
|
||
let bytes = std::fs::read(path).ok()?;
|
||
// NFKC: PDFs often carry CJK radicals ("⽂") for the ordinary characters ("文").
|
||
use unicode_normalization::UnicodeNormalization;
|
||
return std::panic::catch_unwind(|| pdf_extract::extract_text_from_mem(&bytes).ok()).ok().flatten().map(|t| t.nfkc().collect::<String>()).filter(|t| !t.trim().is_empty());
|
||
}
|
||
let bytes = std::fs::read(path).ok()?;
|
||
if bytes.len() > 2 * 1024 * 1024 || bytes.iter().take(8000).any(|b| *b == 0) {
|
||
return None;
|
||
}
|
||
String::from_utf8(bytes).ok()
|
||
}
|
||
|
||
/// What the model gets for a message's attachments: extra text and pictures.
|
||
pub fn for_model(atts: &[Attachment]) -> (String, Vec<String>) {
|
||
let mut text = String::new();
|
||
let mut images = vec![];
|
||
for a in atts {
|
||
let where_ = if a.on_box.is_empty() { format!("檔案存在 AnyBot 所在的電腦：{}", a.path) } else { format!("檔案已放進你自己的電腦：{}", a.on_box) };
|
||
if a.image {
|
||
match image_b64(&a.path) {
|
||
Some(b) => {
|
||
images.push(b);
|
||
text.push_str(&format!("\n\n[附件圖片「{}」，就是下面這張。{where_}]", a.name));
|
||
}
|
||
None => text.push_str(&format!("\n\n[附件圖片「{}」讀不出來。{where_}]", a.name)),
|
||
}
|
||
} else if let Some(t) = read_text(&a.path, &a.name) {
|
||
text.push_str(&format!("\n\n[附件「{}」（{}）。{where_}]\n```\n{}\n```", a.name, human(a.bytes), clip(&t)));
|
||
} else {
|
||
text.push_str(&format!("\n\n[附件「{}」（{}），這裡讀不出內容，可以用工具處理。{where_}]", a.name, human(a.bytes)));
|
||
}
|
||
}
|
||
(text, images)
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn text_and_images_reach_the_model() {
|
||
let dir = std::env::temp_dir().join(format!("anybot-attach-test-{}", std::process::id()));
|
||
std::fs::create_dir_all(&dir).unwrap();
|
||
let txt = dir.join("筆記.md");
|
||
std::fs::write(&txt, "# 標題\n內容 hello").unwrap();
|
||
let png = dir.join("圖.png");
|
||
image::RgbImage::from_pixel(2000, 1000, image::Rgb([10, 20, 30])).save(&png).unwrap();
|
||
let bin = dir.join("data.bin");
|
||
std::fs::write(&bin, [0u8, 1, 2, 0, 5]).unwrap();
|
||
let atts: Vec<Attachment> = [&txt, &png, &bin].iter().map(|p| store(99_999, p).unwrap()).collect();
|
||
assert!(atts[1].image && !atts[0].image);
|
||
let (text, images) = for_model(&atts);
|
||
assert!(text.contains("內容 hello") && text.contains("附件圖片「圖.png」") && text.contains("讀不出內容"), "{text}");
|
||
assert_eq!(images.len(), 1);
|
||
for a in &atts {
|
||
let _ = std::fs::remove_file(&a.path);
|
||
}
|
||
let _ = std::fs::remove_dir_all(&dir);
|
||
assert!(store(1, Path::new("/nonexistent/x")).is_err());
|
||
}
|
||
}
|